From ffb3d2936de6b51bf0334913684c54f7644b20b9 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Thu, 9 Apr 2026 00:03:58 -0700 Subject: [PATCH 01/26] support tvm ffi interfaces Co-Authored-By: rainj-me --- .gitignore | 3 +- CMakeLists.txt | 6 +- build.sh | 4 + csrc/apis/attention.hpp | 4 + csrc/apis/einsum.hpp | 7 +- csrc/apis/gemm.hpp | 4 + csrc/apis/hyperconnection.hpp | 4 + csrc/apis/layout.hpp | 4 + csrc/apis/runtime.hpp | 4 + csrc/jit_kernels/impls/runtime_utils.hpp | 2 +- csrc/jit_kernels/impls/sm100_bf16_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp | 2 +- .../jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp | 2 +- .../impls/sm100_tf32_hc_prenorm_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm90_bf16_gemm.hpp | 2 +- csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp | 2 +- csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp | 2 +- csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp | 2 +- .../impls/sm90_tf32_hc_prenorm_gemm.hpp | 2 +- csrc/jit_kernels/impls/smxx_layout.hpp | 2 +- csrc/tvm_ffi_api.cpp | 467 ++++++++++++++++++ csrc/utils/compatibility.hpp | 2 +- csrc/utils/layout.hpp | 2 +- csrc/utils/math.hpp | 2 +- csrc/utils/torch_compat.hpp | 133 +++++ deep_gemm/__init__.py | 311 +++++++++--- deep_gemm/utils/layout.py | 79 ++- setup.py | 198 ++++---- tests/test_attention.py | 2 +- tests/test_bf16.py | 6 +- tests/test_fp8_fp4.py | 11 +- 31 files changed, 1053 insertions(+), 222 deletions(-) create mode 100644 csrc/tvm_ffi_api.cpp create mode 100644 csrc/utils/torch_compat.hpp diff --git a/.gitignore b/.gitignore index d0cdf6ca42..cb1fcae67d 100644 --- a/.gitignore +++ b/.gitignore @@ -21,4 +21,5 @@ deep_gemm/include/cutlass stubs/ # Symlinks to compiled extensions -deep_gemm/*.so \ No newline at end of file +deep_gemm/*.so +deep_gemm/_C_build \ No newline at end of file diff --git a/CMakeLists.txt b/CMakeLists.txt index bbf625d306..475812a5b9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -5,6 +5,7 @@ set(CMAKE_VERBOSE_MAKEFILE ON) set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -O3 -fPIC -Wno-psabi -Wno-deprecated-declarations") set(CUDA_SEPARABLE_COMPILATION ON) + list(APPEND CUDA_NVCC_FLAGS "-DENABLE_FAST_DEBUG") list(APPEND CUDA_NVCC_FLAGS "-O3") list(APPEND CUDA_NVCC_FLAGS "--ptxas-options=--verbose,--register-usage-level=10,--warn-on-local-memory-usage") @@ -16,13 +17,14 @@ set(TORCH_CUDA_ARCH_LIST "${CUDA_ARCH_LIST}") find_package(CUDAToolkit REQUIRED) find_package(pybind11 REQUIRED) find_package(Torch REQUIRED) +find_package(tvm_ffi REQUIRED) set(CMAKE_CXX_STANDARD 20) set(CMAKE_CUDA_STANDARD 20) include_directories(deep_gemm/include third-party/cutlass/include third-party/cutlass/tools/util/include third-party/fmt/include) -include_directories(${CUDA_TOOLKIT_ROOT_DIR}/include/cccl ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS}) -link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs) +include_directories(${CUDA_TOOLKIT_ROOT_DIR}/targets/x86_64-linux/include ${TORCH_INCLUDE_DIRS} ${PYTHON_INCLUDE_DIRS} ${tvm_ffi_INCLUDE_DIR} ${tvm_ffi_DLPACK_INCLUDE_DIR}) +link_directories(${TORCH_INSTALL_PREFIX}/lib ${CUDA_TOOLKIT_ROOT_DIR}/lib64 ${CUDA_TOOLKIT_ROOT_DIR}/lib64/stubs ${tvm_ffi_ROOT_DIR}/lib) # The main Python API entrance pybind11_add_module(_C csrc/python_api.cpp) diff --git a/build.sh b/build.sh index abdfc40674..90a9deacfc 100755 --- a/build.sh +++ b/build.sh @@ -3,6 +3,10 @@ original_dir=$(pwd) script_dir=$(realpath "$(dirname "$0")") cd "$script_dir" +# Link CUTLASS includes +ln -sf $script_dir/third-party/cutlass/include/cutlass deep_gemm/include +ln -sf $script_dir/third-party/cutlass/include/cute deep_gemm/include + # Remove old dist file, build files, and install rm -rf build dist rm -rf *.egg-info diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 505b0c0973..cecace56ce 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -413,6 +413,8 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q, } #endif +#if 0 + static void register_apis(pybind11::module_& m) { #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE m.def("fp8_gemm_nt_skip_head_mid", &fp8_gemm_nt_skip_head_mid, @@ -450,4 +452,6 @@ static void register_apis(pybind11::module_& m) { #endif } +#endif + } // namespace deep_gemm::attention diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index f82331cacc..8a10ca9074 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -1,8 +1,5 @@ #pragma once -#include -#include - #include "../utils/exception.hpp" #include "../utils/format.hpp" #include "../utils/layout.hpp" @@ -215,6 +212,8 @@ static void fp8_einsum(const std::string& expr, } #endif +#if 0 + static void register_apis(pybind11::module_& m) { #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE m.def("einsum", &einsum, @@ -228,4 +227,6 @@ static void register_apis(pybind11::module_& m) { #endif } +#endif + } // namespace deep_gemm::einsum diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 42622df7d8..41f16a5c3a 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -596,6 +596,8 @@ static void cublaslt_gemm_tt(const torch::Tensor& a, const torch::Tensor& b, cublaslt_gemm_nt(a.transpose(0, 1), b, d, c); } +#if 0 + static void register_apis(pybind11::module_& m) { #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE @@ -712,4 +714,6 @@ static void register_apis(pybind11::module_& m) { py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt); } +#endif + } // namespace deep_gemm::gemm diff --git a/csrc/apis/hyperconnection.hpp b/csrc/apis/hyperconnection.hpp index 1a13984d24..5fc0c002e8 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -59,6 +59,8 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a, #endif +#if 0 + static void register_apis(pybind11::module_& m) { #if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE m.def("tf32_hc_prenorm_gemm", &tf32_hc_prenorm_gemm, @@ -67,4 +69,6 @@ static void register_apis(pybind11::module_& m) { #endif } +#endif + } // namespace deep_gemm::hyperconnection diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index b404241a66..5efb6adb60 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -115,6 +115,8 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te #endif +#if 0 + static void register_apis(pybind11::module_& m) { #if DG_TENSORMAP_COMPATIBLE m.def("transform_sf_into_required_layout", &transform_sf_into_required_layout, @@ -140,4 +142,6 @@ static void register_apis(pybind11::module_& m) { }, py::arg("expected_m") = std::nullopt); } +#endif + } // namespace deep_gemm::layout diff --git a/csrc/apis/runtime.hpp b/csrc/apis/runtime.hpp index 58fef941b7..94cbfd6745 100644 --- a/csrc/apis/runtime.hpp +++ b/csrc/apis/runtime.hpp @@ -8,6 +8,8 @@ namespace deep_gemm::runtime { +#if 0 + static void register_apis(pybind11::module_& m) { m.def("set_num_sms", [&](const int& new_num_sms) { device_runtime->set_num_sms(new_num_sms); @@ -48,4 +50,6 @@ static void register_apis(pybind11::module_& m) { }); } +#endif + } // namespace deep_gemm::runtime diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index 72a76f0d67..02c4ff2e11 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "../heuristics/sm90.hpp" #include "../../jit/handle.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index 26219b0cfd..98cce8f8cb 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp index 65c9d501c2..d2c3e532dc 100644 --- a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp index 07a977d73e..404369a4e9 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp index 0071e2c57f..0665cff094 100644 --- a/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp index 1d29d85585..5c46eb0b88 100644 --- a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp index 473677b70c..78d7f42df7 100644 --- a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp index 9d903d4800..febef94a99 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index 96b5cd0ba4..61b85ec1f0 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp index c17d1b554e..2518079b22 100644 --- a/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index 5d1f17b5b8..7450fba173 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/kernel_runtime.hpp" #include "../../jit/compiler.hpp" diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp new file mode 100644 index 0000000000..e2625f6aa3 --- /dev/null +++ b/csrc/tvm_ffi_api.cpp @@ -0,0 +1,467 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "apis/attention.hpp" +#include "apis/einsum.hpp" +#include "apis/hyperconnection.hpp" +#include "apis/gemm.hpp" +#include "apis/layout.hpp" +#include "tvm/ffi/container/tuple.h" +#include "tvm/ffi/object.h" +#include "utils/torch_compat.hpp" +#include +#include + +using namespace deep_gemm; +using namespace tvm::ffi; + +// --------------------------------------------------------------------------- +// Runtime +// --------------------------------------------------------------------------- +void dg_init(std::string library_root_path, std::string cuda_home) { +#if DG_TENSORMAP_COMPATIBLE + Compiler::prepare_init(library_root_path, cuda_home); + KernelRuntime::prepare_init(cuda_home); +#endif +} + +int64_t dg_get_num_sms() { return device_runtime->get_num_sms(); } +void dg_set_num_sms(int64_t n) { device_runtime->set_num_sms(static_cast(n)); } +int64_t dg_get_compile_mode() { return device_runtime->get_compile_mode(); } +void dg_set_compile_mode(int64_t n) { device_runtime->set_compile_mode(static_cast(n)); } +int64_t dg_get_tc_util() { return device_runtime->get_tc_util(); } +void dg_set_tc_util(int64_t n) { device_runtime->set_tc_util(static_cast(n)); } + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(init, dg_init); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_num_sms, dg_get_num_sms); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_num_sms, dg_set_num_sms); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_compile_mode, dg_get_compile_mode); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_compile_mode, dg_set_compile_mode); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_tc_util, dg_get_tc_util); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_tc_util, dg_set_tc_util); + +// --------------------------------------------------------------------------- +// Layout utilities +// --------------------------------------------------------------------------- +int64_t dg_get_tma_aligned_size(int64_t mn, int64_t element_size) { + return get_tma_aligned_size(static_cast(mn), static_cast(element_size)); +} + +int64_t dg_get_mk_alignment_for_contiguous_layout() { + return get_mk_alignment_for_contiguous_layout(); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_tma_aligned_size, dg_get_tma_aligned_size); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_mk_alignment_for_contiguous_layout, dg_get_mk_alignment_for_contiguous_layout); + +// --------------------------------------------------------------------------- +// Layout kernels +// --------------------------------------------------------------------------- +#if DG_TENSORMAP_COMPATIBLE + +tvm::ffi::Array dg_preprocess_sf(TensorView sf) { + auto sf_v = convert_to_torch_tensor(sf); + auto [dim, ng, mn_pp, sf_k_pp, tma_mn, batched_sf] = preprocess_sf(sf_v); + return {static_cast(dim), + static_cast(ng), + static_cast(mn_pp), + static_cast(sf_k_pp), + static_cast(tma_mn)}; +} + +Tensor dg_get_mn_major_tma_aligned_tensor(TensorView sf) { + auto sf_v = convert_to_torch_tensor(sf); + auto result = get_mn_major_tma_aligned_tensor(sf_v); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +Tensor dg_get_mn_major_tma_aligned_packed_ue8m0_tensor(TensorView sf) { + auto sf_v = convert_to_torch_tensor(sf); + auto result = get_mn_major_tma_aligned_packed_ue8m0_tensor(sf_v); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +Tensor dg_get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + TensorView sf, TensorView ks_tensor, Array ks) { + auto sf_v = convert_to_torch_tensor(sf); + auto ks_tensor_v = convert_to_torch_tensor(ks_tensor); + std::vector ks_opt; + ks_opt.reserve(ks.size()); + + for (Array::iterator it = ks.begin(); it != ks.end(); ++it) { + ks_opt.push_back(static_cast(*it)); + } + auto result = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(sf_v, ks_tensor_v, ks_opt); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +Tensor dg_transform_sf_into_required_layout( + TensorView sf, int64_t mn, int64_t k, + Optional> recipe, + Optional> recipe_ab, + Optional num_groups, + bool is_sfa, + bool disable_ue8m0_cast) { + auto sf_v = convert_to_torch_tensor(sf); + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + auto recipe_ab_opt = recipe_ab.has_value()? std::make_optional(std::make_tuple((int)recipe_ab.value().get<0>(), (int)recipe_ab.value().get<1>())) : std::nullopt; + std::optional ng = num_groups.has_value() ? std::make_optional(static_cast(num_groups.value())) : std::nullopt; + auto result = layout::transform_sf_into_required_layout( + sf_v, static_cast(mn), static_cast(k), + recipe_opt, recipe_ab_opt, ng, is_sfa, disable_ue8m0_cast); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(preprocess_sf, dg_preprocess_sf); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_mn_major_tma_aligned_tensor, dg_get_mn_major_tma_aligned_tensor); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_mn_major_tma_aligned_packed_ue8m0_tensor, dg_get_mn_major_tma_aligned_packed_ue8m0_tensor); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor, dg_get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(transform_sf_into_required_layout, dg_transform_sf_into_required_layout); + +#endif // DG_TENSORMAP_COMPATIBLE + +// --------------------------------------------------------------------------- +// cuBLASLt GEMMs (always available) +// --------------------------------------------------------------------------- +void dg_cublaslt_gemm_nt(TensorView a, TensorView b, TensorView d, Optional c) { + auto c_val = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::cublaslt_gemm_nt( + convert_to_torch_tensor(a), + convert_to_torch_tensor(b), + convert_to_torch_tensor(d), + c_val + ); +} +void dg_cublaslt_gemm_nn(TensorView a, TensorView b, TensorView d, Optional c) { + auto c_val = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::cublaslt_gemm_nn( + convert_to_torch_tensor(a), + convert_to_torch_tensor(b), + convert_to_torch_tensor(d), + c_val + ); +} +void dg_cublaslt_gemm_tn(TensorView a, TensorView b, TensorView d, Optional c) { + auto c_val = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::cublaslt_gemm_tn( + convert_to_torch_tensor(a), + convert_to_torch_tensor(b), + convert_to_torch_tensor(d), + c_val + ); +} +void dg_cublaslt_gemm_tt(TensorView a, TensorView b, TensorView d, Optional c) { + auto c_val = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::cublaslt_gemm_tt( + convert_to_torch_tensor(a), + convert_to_torch_tensor(b), + convert_to_torch_tensor(d), + c_val + ); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(cublaslt_gemm_nt, dg_cublaslt_gemm_nt); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(cublaslt_gemm_nn, dg_cublaslt_gemm_nn); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(cublaslt_gemm_tn, dg_cublaslt_gemm_tn); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(cublaslt_gemm_tt, dg_cublaslt_gemm_tt); + +// --------------------------------------------------------------------------- +// FP8/FP4 GEMMs and BF16 GEMMs +// --------------------------------------------------------------------------- +#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE + +void dg_fp8_fp4_gemm_nt(TensorView a, TensorView a_sf, + TensorView b, TensorView b_sf, + TensorView d, + Optional c, + Optional> recipe, + Optional> recipe_a, + Optional> recipe_b, + std::string compiled_dims, + bool disable_ue8m0_cast) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; + auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; + gemm::fp8_fp4_gemm_nt(std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), c_opt, + recipe_opt, recipe_a_opt, recipe_b_opt, + compiled_dims, disable_ue8m0_cast); +} + +void dg_fp8_fp4_gemm_nn(TensorView a, TensorView a_sf, + TensorView b, TensorView b_sf, + TensorView d, + Optional c, + Optional> recipe, + Optional> recipe_a, + Optional> recipe_b, + std::string compiled_dims, + bool disable_ue8m0_cast) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; + auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; + gemm::fp8_fp4_gemm_nn(std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), c_opt, + recipe_opt, recipe_a_opt, recipe_b_opt, + compiled_dims, disable_ue8m0_cast); +} + +void dg_fp8_fp4_gemm_tn(TensorView a, TensorView a_sf, + TensorView b, TensorView b_sf, + TensorView d, + Optional c, + Optional> recipe, + Optional> recipe_a, + Optional> recipe_b, + std::string compiled_dims, + bool disable_ue8m0_cast) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; + auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; + gemm::fp8_fp4_gemm_tn(std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), c_opt, + recipe_opt, recipe_a_opt, recipe_b_opt, + compiled_dims, disable_ue8m0_cast); +} + +void dg_fp8_fp4_gemm_tt(TensorView a, TensorView a_sf, + TensorView b, TensorView b_sf, + TensorView d, + Optional c, + Optional> recipe, + Optional> recipe_a, + Optional> recipe_b, + std::string compiled_dims, + bool disable_ue8m0_cast) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; + auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; + gemm::fp8_fp4_gemm_tt(std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), c_opt, + recipe_opt, recipe_a_opt, recipe_b_opt, + compiled_dims, disable_ue8m0_cast); +} + +void dg_m_grouped_fp8_fp4_gemm_nt_contiguous(TensorView a, TensorView a_sf, + TensorView b, TensorView b_sf, + TensorView d, + TensorView grouped_layout, + Optional> recipe, + Optional> recipe_a, + Optional> recipe_b, + std::string compiled_dims, + bool disable_ue8m0_cast, + bool use_psum_layout, + Optional expected_m_for_psum_layout) { + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; + auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; + auto expected_m_for_psum_layout_opts = expected_m_for_psum_layout.has_value()? std::make_optional((int) expected_m_for_psum_layout.value()) : std::nullopt; + gemm::m_grouped_fp8_fp4_gemm_nt_contiguous( + std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), convert_to_torch_tensor(grouped_layout), + recipe_opt, recipe_a_opt, recipe_b_opt, + compiled_dims, disable_ue8m0_cast, + use_psum_layout, expected_m_for_psum_layout_opts + ); +} + +void dg_m_grouped_fp8_fp4_gemm_nn_contiguous(TensorView a, TensorView a_sf, + TensorView b, TensorView b_sf, + TensorView d, + TensorView grouped_layout, + Optional> recipe, + Optional> recipe_a, + Optional> recipe_b, + std::string compiled_dims, + bool disable_ue8m0_cast, + bool use_psum_layout) { + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; + auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; + gemm::m_grouped_fp8_fp4_gemm_nn_contiguous( + std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), convert_to_torch_tensor(grouped_layout), + recipe_opt, recipe_a_opt, recipe_b_opt, + compiled_dims, disable_ue8m0_cast, use_psum_layout + ); +} + +Tuple, Optional> +dg_m_grouped_fp8_fp4_gemm_nt_masked(TensorView a, TensorView a_sf, + TensorView b, TensorView b_sf, + TensorView d, + TensorView masked_m, + int64_t expected_m, + Optional> recipe, + Optional> recipe_a, + Optional> recipe_b, + std::string compiled_dims, + bool disable_ue8m0_cast, + int64_t max_block_n, + bool enable_overlap, + Optional signal) { + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; + auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; + auto signal_opt = signal.has_value() ? std::make_optional(convert_to_torch_tensor(signal.value())) : std::nullopt; + auto result = gemm::m_grouped_fp8_fp4_gemm_nt_masked( + std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), convert_to_torch_tensor(masked_m), + (int) expected_m, recipe_opt, recipe_a_opt, recipe_b_opt, + compiled_dims, disable_ue8m0_cast, (int) max_block_n, enable_overlap, signal_opt + ); + auto [first, second] = result; + Optional first_res, second_res; + if(first.has_value()) { + first_res = (int64_t) first.value(); + } + if (second.has_value()) { + second_res = (int64_t) second.value(); + } + Tuple, Optional> res(first_res, second_res); + return res; +} + + +void dg_bf16_gemm_nt(TensorView a, TensorView b, TensorView d, + Optional c, + std::string compiled_dims) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::bf16_gemm_nt(convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, compiled_dims); +} +void dg_bf16_gemm_nn(TensorView a, TensorView b, TensorView d, + Optional c, + std::string compiled_dims) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::bf16_gemm_nn(convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, compiled_dims); +} +void dg_bf16_gemm_tn(TensorView a, TensorView b, TensorView d, + Optional c, + std::string compiled_dims) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::bf16_gemm_tn(convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, compiled_dims); +} +void dg_bf16_gemm_tt(TensorView a, TensorView b, TensorView d, + Optional c, + std::string compiled_dims) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::bf16_gemm_tt(convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, compiled_dims); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_nt, dg_fp8_fp4_gemm_nt); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_nn, dg_fp8_fp4_gemm_nn); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_tn, dg_fp8_fp4_gemm_tn); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_tt, dg_fp8_fp4_gemm_tt); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_fp8_fp4_gemm_nt_contiguous, dg_m_grouped_fp8_fp4_gemm_nt_contiguous); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_fp8_fp4_gemm_nt_masked, dg_m_grouped_fp8_fp4_gemm_nt_masked); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_nt, dg_bf16_gemm_nt); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_nn, dg_bf16_gemm_nn); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_tn, dg_bf16_gemm_tn); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_tt, dg_bf16_gemm_tt); + +// Einsum +void dg_einsum(std::string expr, TensorView a, TensorView b, TensorView d, + Optional c, bool use_cublaslt) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + einsum::einsum(expr, convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, use_cublaslt); +} + +void dg_fp8_einsum(std::string expr, + TensorView a_data, TensorView a_sf, + TensorView b_data, TensorView b_sf, + TensorView d, + Optional c, + Tuple recipe) { + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + auto recipe_opt = std::make_tuple((int)recipe.get<0>(), (int)recipe.get<1>(), (int)recipe.get<2>()); + einsum::fp8_einsum(expr, std::make_pair(convert_to_torch_tensor(a_data), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b_data), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), c_opt, recipe_opt); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(einsum, dg_einsum); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_einsum, dg_fp8_einsum); + +// Hyperconnection +void dg_tf32_hc_prenorm_gemm(TensorView a, TensorView b, TensorView d, + TensorView sqr_sum, Optional num_splits) { + auto ns = num_splits.has_value() ? std::make_optional(static_cast(num_splits.value())) : std::nullopt; + hyperconnection::tf32_hc_prenorm_gemm(convert_to_torch_tensor(a), convert_to_torch_tensor(b), + convert_to_torch_tensor(d), convert_to_torch_tensor(sqr_sum), ns); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(tf32_hc_prenorm_gemm, dg_tf32_hc_prenorm_gemm); + +// Attention +void dg_fp8_gemm_nt_skip_head_mid(TensorView a_data, TensorView a_sf, + TensorView b_data, TensorView b_sf, + TensorView d, + Tuple head_splits, + Optional> recipe, + std::string compiled_dims, + bool disable_ue8m0_cast) { + auto head_splits_opt = std::make_tuple((int)head_splits.get<0>(), (int)head_splits.get<1>(), (int)head_splits.get<2>()); + auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; + attention::fp8_gemm_nt_skip_head_mid( + std::make_pair(convert_to_torch_tensor(a_data), convert_to_torch_tensor(a_sf)), + std::make_pair(convert_to_torch_tensor(b_data), convert_to_torch_tensor(b_sf)), + convert_to_torch_tensor(d), head_splits_opt, recipe_opt, compiled_dims, disable_ue8m0_cast); +} + +Tensor dg_fp8_mqa_logits(TensorView q, TensorView kv_data, TensorView kv_sf, + TensorView weights, TensorView cu_seq_len_k_start, + TensorView cu_seq_len_k_end, + bool clean_logits, int64_t max_seqlen_k) { + auto result = attention::fp8_mqa_logits( + convert_to_torch_tensor(q), + std::make_pair(convert_to_torch_tensor(kv_data), convert_to_torch_tensor(kv_sf)), + convert_to_torch_tensor(weights), convert_to_torch_tensor(cu_seq_len_k_start), + convert_to_torch_tensor(cu_seq_len_k_end), clean_logits, static_cast(max_seqlen_k)); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +Tensor dg_get_paged_mqa_logits_metadata(TensorView context_lens, int64_t block_kv, + int64_t num_sms) { + auto result = attention::get_paged_mqa_logits_metadata( + convert_to_torch_tensor(context_lens), static_cast(block_kv), + static_cast(num_sms)); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +Tensor dg_fp8_paged_mqa_logits(TensorView q, TensorView fused_kv_cache, + TensorView weights, TensorView context_lens, + TensorView block_table, TensorView schedule_meta, + int64_t max_context_len, bool clean_logits) { + auto result = attention::fp8_paged_mqa_logits( + convert_to_torch_tensor(q), convert_to_torch_tensor(fused_kv_cache), + convert_to_torch_tensor(weights), convert_to_torch_tensor(context_lens), + convert_to_torch_tensor(block_table), convert_to_torch_tensor(schedule_meta), + static_cast(max_context_len), clean_logits); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_gemm_nt_skip_head_mid, dg_fp8_gemm_nt_skip_head_mid); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_mqa_logits, dg_fp8_mqa_logits); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_paged_mqa_logits_metadata, dg_get_paged_mqa_logits_metadata); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_paged_mqa_logits, dg_fp8_paged_mqa_logits); + +#endif // DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE diff --git a/csrc/utils/compatibility.hpp b/csrc/utils/compatibility.hpp index 9e2d67205a..cc0cac2a60 100644 --- a/csrc/utils/compatibility.hpp +++ b/csrc/utils/compatibility.hpp @@ -1,8 +1,8 @@ #pragma once -#include #include #include +#include // `torch::kFloat8_e4m3fn` is supported since PyTorch 2.1 #define DG_FP8_COMPATIBLE (TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 1)) diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 928472d35a..62dc55e586 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +#include #include "math.hpp" #include "exception.hpp" diff --git a/csrc/utils/math.hpp b/csrc/utils/math.hpp index 19a86c3890..f999868c08 100644 --- a/csrc/utils/math.hpp +++ b/csrc/utils/math.hpp @@ -1,7 +1,7 @@ // TODO: merge this file with `math.cuh` (the device part) #pragma once -#include +#include #include "exception.hpp" diff --git a/csrc/utils/torch_compat.hpp b/csrc/utils/torch_compat.hpp new file mode 100644 index 0000000000..d1a02443a0 --- /dev/null +++ b/csrc/utils/torch_compat.hpp @@ -0,0 +1,133 @@ +#pragma once + +#include + +#include +#include +#include + +#include +#include +#include +#include + +#include "exception.hpp" + +namespace deep_gemm { + +// --------------------------------------------------------------------------- +// Scalar-type → CUDA / cuBLAS / TensorMap conversions +// --------------------------------------------------------------------------- +inline std::string dtype_to_cuda_type_string(at::ScalarType dtype) { + if (dtype == at::kInt) return "int"; + if (dtype == at::kFloat) return "float"; + if (dtype == at::kBFloat16) return "cutlass::bfloat16_t"; + if (dtype == at::kFloat8_e4m3fn) return "cutlass::float_e4m3_t"; + if (dtype == at::kByte) return "cutlass::detail::float_e2m1_unpacksmem_t"; + DG_HOST_UNREACHABLE("Unsupported dtype for CUDA type string"); +} + +inline cudaDataType_t dtype_to_cublas(at::ScalarType dtype) { + if (dtype == at::kFloat) return CUDA_R_32F; + if (dtype == at::kHalf) return CUDA_R_16F; + if (dtype == at::kBFloat16) return CUDA_R_16BF; + if (dtype == at::kFloat8_e4m3fn) return CUDA_R_8F_E4M3; + if (dtype == at::kInt) return CUDA_R_32I; + DG_HOST_UNREACHABLE("Unsupported dtype for cuBLAS"); +} + +inline CUtensorMapDataType dtype_to_tensormap(at::ScalarType dtype, bool allow_tf32 = false) { + if (allow_tf32 && dtype == at::kFloat) return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32; + if (dtype == at::kInt) return CU_TENSOR_MAP_DATA_TYPE_INT32; + if (dtype == at::kFloat) return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; + if (dtype == at::kBFloat16) return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; + if (dtype == at::kFloat8_e4m3fn) return CU_TENSOR_MAP_DATA_TYPE_UINT8; + if (dtype == at::kByte) return CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B; + DG_HOST_UNREACHABLE("Unsupported dtype for TensorMap"); +} + +inline std::string dtype_to_string(at::ScalarType dtype) { + return std::string(c10::toString(dtype)); +} + +// --------------------------------------------------------------------------- +// CUDA stream helper +// --------------------------------------------------------------------------- +inline cudaStream_t get_current_cuda_stream() { + return at::cuda::getCurrentCUDAStream().stream(); +} + +// --------------------------------------------------------------------------- +// Create torch::Tensor from raw device pointer (non-owning view) +// --------------------------------------------------------------------------- +inline torch::Tensor tensor_from_ptr(void* data, at::ScalarType dtype, + std::initializer_list shape, + int device_id = 0) { + auto opts = torch::TensorOptions().dtype(dtype) + .device(torch::kCUDA, device_id) + .requires_grad(false); + auto sizes = std::vector(shape); + return torch::from_blob(data, sizes, opts); +} + +inline torch::Tensor tensor_from_ptr_strided(void* data, at::ScalarType dtype, + std::initializer_list shape, + std::initializer_list strides, + int device_id = 0) { + auto opts = torch::TensorOptions().dtype(dtype) + .device(torch::kCUDA, device_id) + .requires_grad(false); + auto sizes = std::vector(shape); + auto str = std::vector(strides); + return torch::from_blob(data, sizes, str, opts); +} + +// --------------------------------------------------------------------------- +// DLTensor* → torch::Tensor (non-owning, for tvm-ffi boundary) +// --------------------------------------------------------------------------- +inline at::ScalarType dl_dtype_to_torch(DLDataType dtype) { + if (dtype.lanes != 1) DG_HOST_UNREACHABLE("Unsupported DLDataType lanes"); + switch (dtype.code) { + case kDLFloat: + if (dtype.bits == 64) return at::kDouble; + if (dtype.bits == 32) return at::kFloat; + if (dtype.bits == 16) return at::kHalf; + break; + case kDLBfloat: + if (dtype.bits == 16) return at::kBFloat16; + break; + case kDLInt: + if (dtype.bits == 64) return at::kLong; + if (dtype.bits == 32) return at::kInt; + if (dtype.bits == 16) return at::kShort; + if (dtype.bits == 8) return at::kChar; + break; + case kDLUInt: + if (dtype.bits == 8) return at::kByte; + break; + case kDLFloat8_e4m3fn: // kDLFloat8_e4m3fn + return at::kFloat8_e4m3fn; + default: + break; + } + DG_HOST_UNREACHABLE("Unsupported DLDataType for torch conversion: " + tvm::ffi::DLDataTypeToString(dtype)); +} + +inline torch::Tensor convert_to_torch_tensor(tvm::ffi::TensorView tensor) { + auto scalar_type = dl_dtype_to_torch(tensor.dtype()); + int device_id = tensor.device().device_id; + void* data = static_cast(tensor.data_ptr()) + tensor.byte_offset(); + + auto sizes = std::vector(tensor.shape().begin(), tensor.shape().end()); + auto opts = torch::TensorOptions().dtype(scalar_type) + .device(torch::kCUDA, device_id) + .requires_grad(false); + + if (tensor.strides().data()) { + auto strides = std::vector(tensor.strides().begin(), tensor.strides().end()); + return torch::from_blob(data, sizes, strides, opts); + } + return torch::from_blob(data, sizes, opts); +} + +} // namespace deep_gemm diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index a9542e2f44..9201b73ef5 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -1,10 +1,18 @@ +from __future__ import annotations + import os +import shutil import subprocess +from typing import TYPE_CHECKING + import torch +import tvm_ffi + +if TYPE_CHECKING: + from tvm_ffi.module import Module # Set some default environment provided at setup try: - # noinspection PyUnresolvedReferences from .envs import persistent_envs for key, value in persistent_envs.items(): if key not in os.environ: @@ -12,72 +20,229 @@ except ImportError: pass -# Configs -from . import _C -from ._C import ( - set_num_sms, - get_num_sms, - set_tc_util, - get_tc_util, - set_ignore_compile_dims, - set_block_size_multiple_of, - set_pdl, - get_pdl, -) +# --------------------------------------------------------------------------- +# Build & load the tvm-ffi _C module +# --------------------------------------------------------------------------- +def _find_cuda_home() -> str: + cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') + if cuda_home is None: + try: + with open(os.devnull, 'w') as devnull: + nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n') + cuda_home = os.path.dirname(os.path.dirname(nvcc)) + except Exception: + cuda_home = '/usr/local/cuda' + if not os.path.exists(cuda_home): + cuda_home = None + assert cuda_home is not None + return cuda_home + + +def _build_module(pkg_dir: str, cuda_home: str) -> str: + """Build the _C shared library using tvm_ffi.cpp.build().""" + import tvm_ffi.cpp + + root_dir = os.path.dirname(pkg_dir) + cxx_abi = int(torch.compiled_with_cxx11_abi()) + + os.environ.setdefault('TVM_FFI_CUDA_ARCH_LIST', _get_cuda_arch()) + + extra_cflags = [ + '-std=c++17', '-O3', '-fPIC', + '-Wno-psabi', '-Wno-deprecated-declarations', + f'-D_GLIBCXX_USE_CXX11_ABI={cxx_abi}', + ] + if int(os.environ.get('DG_JIT_USE_RUNTIME_API', '0')): + extra_cflags.append('-DDG_JIT_USE_RUNTIME_API') + + # Torch include/lib paths + torch_dir = os.path.dirname(torch.__file__) + torch_include = os.path.join(torch_dir, 'include') + torch_include_csrc = os.path.join(torch_include, 'torch', 'csrc', 'api', 'include') + torch_lib = os.path.join(torch_dir, 'lib') + + import sysconfig + python_include = sysconfig.get_path('include') + + extra_include_paths = [ + f'{cuda_home}/include', + python_include, + torch_include, + torch_include_csrc, + os.path.join(root_dir, 'deep_gemm', 'include'), + os.path.join(root_dir, 'third-party', 'cutlass', 'include'), + os.path.join(root_dir, 'third-party', 'fmt', 'include'), + ] + cccl_path = f'{cuda_home}/include/cccl' + if os.path.exists(cccl_path): + extra_include_paths.append(cccl_path) + + extra_ldflags = [ + f'-L{cuda_home}/lib64', + f'-L{torch_lib}', + '-lcudart', + '-lnvrtc', + '-lcublasLt', + '-lcublas', + '-ltorch', + '-ltorch_cpu', + '-lc10', + '-lc10_cuda', + '-ltorch_cuda', + ] + + build_dir = os.path.join(pkg_dir, '_C_build') + os.makedirs(build_dir, exist_ok=True) + + lib_path = tvm_ffi.cpp.build( + name='_C', + cpp_files=[os.path.join(root_dir, 'csrc', 'tvm_ffi_api.cpp')], + extra_cflags=extra_cflags, + extra_ldflags=extra_ldflags, + extra_include_paths=extra_include_paths, + build_directory=build_dir, + ) + # Copy the .so into the package directory for easy loading + target = os.path.join(pkg_dir, '_C.so') + shutil.copy2(lib_path, target) + return target + + +def _get_cuda_arch() -> str: + try: + status = subprocess.run( + args=['nvidia-smi', '--query-gpu=compute_cap', '--format=csv,noheader'], + capture_output=True, check=True, + ) + return status.stdout.decode('utf-8').strip().split('\n')[0] + except Exception: + return '9.0' + + +def _load_module() -> Module: + """Load (or build then load) the compiled tvm-ffi module.""" + pkg_dir = os.path.dirname(os.path.abspath(__file__)) + lib_path = os.path.join(pkg_dir, '_C.so') + + if not os.path.exists(lib_path): + cuda_home = _find_cuda_home() + print(f'[DeepGEMM] Building _C module with tvm-ffi (CUDA_HOME={cuda_home})...') + lib_path = _build_module(pkg_dir, cuda_home) + print(f'[DeepGEMM] Built _C module: {lib_path}') + + return tvm_ffi.load_module(lib_path) + +_C: Module = _load_module() + +# --------------------------------------------------------------------------- +# Runtime config +# --------------------------------------------------------------------------- +set_num_sms = _C.set_num_sms +get_num_sms = _C.get_num_sms +set_compile_mode = _C.set_compile_mode +get_compile_mode = _C.get_compile_mode +set_tc_util = _C.set_tc_util +get_tc_util = _C.get_tc_util # cuBLASLt Kernels -from ._C import ( - cublaslt_gemm_nt, cublaslt_gemm_nn, - cublaslt_gemm_tn, cublaslt_gemm_tt, -) +cublaslt_gemm_nt = _C.cublaslt_gemm_nt +cublaslt_gemm_nn = _C.cublaslt_gemm_nn +cublaslt_gemm_tn = _C.cublaslt_gemm_tn +cublaslt_gemm_tt = _C.cublaslt_gemm_tt + +def _parse_tensor_or_tuple(input): + if type(input) is tuple or type(input) is list: + return input[0], input[1] + elif isinstance(input, torch.Tensor): + scale = torch.Tensor([1.0], dtype=torch.float32, device=input.device) + return input, scale + assert False, "Expected Tensor, (Tensor, Tensor) tuple, or [Tensor, Tensor] list" + +# --------------------------------------------------------------------------- +# GEMM / Attention / Einsum wrappers (handle optional params in Python) +# --------------------------------------------------------------------------- try: - # DeepGEMM Kernels - from ._C import ( - # FP8 FP4 GEMMs - fp8_fp4_gemm_nt, fp8_fp4_gemm_nn, - fp8_fp4_gemm_tn, fp8_fp4_gemm_tt, - m_grouped_fp8_fp4_gemm_nt_contiguous, - m_grouped_fp8_fp4_gemm_nn_contiguous, - m_grouped_fp8_fp4_gemm_nt_masked, - # FP8 GEMMs - fp8_gemm_nt, fp8_gemm_nn, - fp8_gemm_tn, fp8_gemm_tt, - fp8_gemm_nt_skip_head_mid, - m_grouped_fp8_gemm_nt_contiguous, - m_grouped_fp8_gemm_nn_contiguous, - m_grouped_fp8_gemm_nt_masked, - k_grouped_fp8_gemm_nt_contiguous, - k_grouped_fp8_gemm_tn_contiguous, - # BF16 GEMMs - bf16_gemm_nt, bf16_gemm_nn, - bf16_gemm_tn, bf16_gemm_tt, - m_grouped_bf16_gemm_nt_contiguous, - m_grouped_bf16_gemm_nn_contiguous, - m_grouped_bf16_gemm_nt_masked, - k_grouped_bf16_gemm_tn_contiguous, - # Einsum kernels - einsum, - fp8_einsum, - # Attention kernels - fp8_fp4_mqa_logits, - get_paged_mqa_logits_metadata, - fp8_fp4_paged_mqa_logits, - # Attention kernels (legacy) - fp8_mqa_logits, - fp8_paged_mqa_logits, - # Hyperconnection kernels - tf32_hc_prenorm_gemm, - # Layout kernels - transform_sf_into_required_layout, - ) + def fp8_fp4_gemm_nt(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_fp4_gemm_nt(a_data, a_sf, b_data, b_sf, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) - # Some alias for legacy supports - # TODO: remove these later - fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_gemm_nt_masked - bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked -except ImportError: - # Expected behavior for CUDA runtime version before 12.1 + def fp8_fp4_gemm_nn(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_fp4_gemm_nn(a_data, a_sf, b_data, b_sf, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + + def fp8_fp4_gemm_tn(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_fp4_gemm_tn(a_data, a_sf, b_data, b_sf, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + + def fp8_fp4_gemm_tt(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_fp4_gemm_tt(a_data, a_sf, b_data, b_sf, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + + fp8_gemm_nt = fp8_fp4_gemm_nt + fp8_gemm_nn = fp8_fp4_gemm_nn + fp8_gemm_tn = fp8_fp4_gemm_tn + fp8_gemm_tt = fp8_fp4_gemm_tt + + def m_grouped_fp8_fp4_gemm_nt_contiguous(a, b, d, grouped_layout, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False, use_psum_layout=False, expected_m_for_psum_layout=None): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.m_grouped_fp8_fp4_gemm_nt_contiguous(a_data, a_sf, b_data, b_sf, d, grouped_layout, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, use_psum_layout, expected_m_for_psum_layout) + + def m_grouped_fp8_fp4_gemm_nn_contiguous(a, b, d, grouped_layout, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False, use_psum_layout=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.m_grouped_fp8_fp4_gemm_nn_contiguous(a_data, a_sf, b_data, b_sf, d, grouped_layout, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, use_psum_layout) + + m_grouped_fp8_gemm_nt_contiguous = m_grouped_fp8_fp4_gemm_nt_contiguous + m_grouped_fp8_gemm_nn_contiguous = m_grouped_fp8_fp4_gemm_nn_contiguous + + def bf16_gemm_nt(a, b, d, c=None, compiled_dims=''): + _C.bf16_gemm_nt(a, b, d, c, compiled_dims) + + def bf16_gemm_nn(a, b, d, c=None, compiled_dims=''): + _C.bf16_gemm_nn(a, b, d, c, compiled_dims) + + def bf16_gemm_tn(a, b, d, c=None, compiled_dims=''): + _C.bf16_gemm_tn(a, b, d, c, compiled_dims) + + def bf16_gemm_tt(a, b, d, c=None, compiled_dims=''): + _C.bf16_gemm_tt(a, b, d, c, compiled_dims) + + def einsum(expr, a, b, d, c=None, use_cublaslt=False): + _C.einsum(expr, a, b, d, c, use_cublaslt) + + def fp8_einsum(expr, a, b, d, c=None, recipe=(1, 128, 128)): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_einsum(expr, a_data, a_sf, b_data, b_sf, d, c, recipe) + + def fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, recipe=None, compiled_dims='nk', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_gemm_nt_skip_head_mid(a_data, a_sf, b_data, b_sf, d, head_splits, recipe, compiled_dims, disable_ue8m0_cast) + + def fp8_paged_mqa_logits(q, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits=False): + return _C.fp8_paged_mqa_logits(q, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits) + + def fp8_mqa_logits(q, kv_data, kv_sf, weights, ks, ke, clean_logits=False, max_seqlen_k=0): + return _C.fp8_mqa_logits(q, kv_data, kv_sf, weights, ks, ke, clean_logits, max_seqlen_k) + + def get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms): + return _C.get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms) + + def tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits=None): + _C.tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits) + + def transform_sf_into_required_layout(sf, mn, k, recipe, recipe_ab = None, num_groups=None, is_sfa=False, disable_ue8m0_cast=False): + return _C.transform_sf_into_required_layout(sf, mn, k, recipe, recipe_ab, num_groups, is_sfa, disable_ue8m0_cast) + + get_mk_alignment_for_contiguous_layout = _C.get_mk_alignment_for_contiguous_layout + + def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False, max_block_n=256, enable_overlap=False, signal=None): + (a, a_sf), (b, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, max_block_n, enable_overlap, signal) + + fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_fp4_gemm_nt_masked + bf16_m_grouped_gemm_nt_masked = None + +except AttributeError: pass # Mega kernels @@ -100,27 +265,9 @@ print(f'Failed to load legacy DeepGEMM A100 Triton kernels: {e}') # Initialize CPP modules -def _find_cuda_home() -> str: - # TODO: reuse PyTorch API later - # For some PyTorch versions, the original `_find_cuda_home` will initialize CUDA, which is incompatible with process forks - cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') - if cuda_home is None: - # noinspection PyBroadException - try: - with open(os.devnull, 'w') as devnull: - nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n') - cuda_home = os.path.dirname(os.path.dirname(nvcc)) - except Exception: - cuda_home = '/usr/local/cuda' - if not os.path.exists(cuda_home): - cuda_home = None - assert cuda_home is not None - return cuda_home - - _C.init( - os.path.dirname(os.path.abspath(__file__)), # Library root directory path - _find_cuda_home() # CUDA home + os.path.dirname(os.path.abspath(__file__)), + _find_cuda_home() ) __version__ = '2.5.0' diff --git a/deep_gemm/utils/layout.py b/deep_gemm/utils/layout.py index 6512c5ab7a..7dc21cb788 100644 --- a/deep_gemm/utils/layout.py +++ b/deep_gemm/utils/layout.py @@ -1,21 +1,66 @@ -try: - from .._C import ( - get_tma_aligned_size, - get_mn_major_tma_aligned_tensor, - get_mn_major_tma_aligned_packed_ue8m0_tensor, - get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor - ) -except ImportError: - # Expected behavior for CUDA runtime version before 12.1 - pass +"""Layout utility wrappers. + +These functions call into the tvm-ffi _C module for CUDA kernel execution +and handle output tensor allocation on the Python side. +""" +from __future__ import annotations + +import torch +from .math import align, ceil_div + -# Valid for all CUDA versions -from .._C import ( - set_mk_alignment_for_contiguous_layout, - get_mk_alignment_for_contiguous_layout, - get_theoretical_mk_alignment_for_contiguous_layout, -) +def _get_C(): + from .. import _C + return _C + + +def get_tma_aligned_size(mn: int, element_size: int) -> int: + return int(_get_C().get_tma_aligned_size(mn, element_size)) + + +def get_mk_alignment_for_contiguous_layout() -> int: + return int(_get_C().get_mk_alignment_for_contiguous_layout()) -# Some alias get_m_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout get_k_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout + + +def _pack_fp32_into_ue8m0_fallback(x: torch.Tensor) -> torch.Tensor: + """PyTorch fallback for ue8m0 packing when the CUDA kernel cannot handle the layout.""" + assert x.dtype == torch.float and x.dim() in (2, 3) + mn, k = x.shape[-2], x.shape[-1] + remove_dim = False + if x.dim() == 2: + x, remove_dim = x.unsqueeze(0), True + b = x.shape[0] + ue8m0_tensor = (x.view(torch.int) >> 23).to(torch.uint8) + aligned_mn = get_tma_aligned_size(mn, 4) + aligned_k = align(k, 4) + padded = torch.zeros((b, aligned_mn, aligned_k), device=x.device, dtype=torch.uint8) + padded[:, :mn, :k] = ue8m0_tensor + padded = padded.view(-1).view(dtype=torch.int).view(b, aligned_mn, aligned_k // 4) + transposed = torch.zeros((b, aligned_k // 4, aligned_mn), device=x.device, dtype=torch.int).mT + transposed[:, :, :] = padded + aligned_x = transposed[:, :mn, :] + return aligned_x.squeeze(0) if remove_dim else aligned_x + + +try: + _get_C().preprocess_sf # probe availability + + def get_mn_major_tma_aligned_tensor(sf: torch.Tensor) -> torch.Tensor: + """Transpose FP32 scaling factors into MN-major TMA-aligned layout.""" + return _get_C().get_mn_major_tma_aligned_tensor(sf) + + def get_mn_major_tma_aligned_packed_ue8m0_tensor(sf: torch.Tensor) -> torch.Tensor: + """Pack FP32 scaling factors into UE8M0 int32 in MN-major TMA-aligned layout.""" + return _get_C().get_mn_major_tma_aligned_packed_ue8m0_tensor(sf) + + def get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + sf: torch.Tensor, ks_tensor: torch.Tensor, ks: list[int]) -> torch.Tensor: + """Pack k-grouped FP32 scaling factors into UE8M0 int32 in MN-major TMA-aligned layout.""" + return _get_C().get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + sf, ks_tensor, ks) + +except AttributeError: + pass diff --git a/setup.py b/setup.py index c4d74ae929..b114487445 100644 --- a/setup.py +++ b/setup.py @@ -5,50 +5,40 @@ import setuptools import subprocess import sys -import torch import platform -import urllib -import urllib.error -import urllib.request from setuptools import find_packages from setuptools.command.build_py import build_py from packaging.version import parse from pathlib import Path -from torch.utils.cpp_extension import CUDAExtension, CUDA_HOME -from wheel.bdist_wheel import bdist_wheel as _bdist_wheel -from scripts.generate_pyi import generate_pyi_file - DG_SKIP_CUDA_BUILD = int(os.getenv('DG_SKIP_CUDA_BUILD', '0')) == 1 DG_FORCE_BUILD = int(os.getenv('DG_FORCE_BUILD', '0')) == 1 DG_USE_LOCAL_VERSION = int(os.getenv('DG_USE_LOCAL_VERSION', '1')) == 1 DG_JIT_USE_RUNTIME_API = int(os.environ.get('DG_JIT_USE_RUNTIME_API', '0')) == 1 -# Compiler flags -cxx_flags = ['-std=c++17', '-O3', '-fPIC', '-Wno-psabi', '-Wno-deprecated-declarations', - f'-D_GLIBCXX_USE_CXX11_ABI={int(torch.compiled_with_cxx11_abi())}'] -if DG_JIT_USE_RUNTIME_API: - cxx_flags.append('-DDG_JIT_USE_RUNTIME_API') - -# Sources current_dir = os.path.dirname(os.path.realpath(__file__)) -sources = ['csrc/python_api.cpp'] -build_include_dirs = [ - f'{CUDA_HOME}/include', - f'{CUDA_HOME}/include/cccl', - 'deep_gemm/include', - 'third-party/cutlass/include', - 'third-party/fmt/include', -] -build_libraries = ['cudart', 'nvrtc'] -build_library_dirs = [f'{CUDA_HOME}/lib64'] + third_party_include_dirs = [ 'third-party/cutlass/include/cute', 'third-party/cutlass/include/cutlass', ] -# Release -base_wheel_url = 'https://github.com/DeepSeek-AI/DeepGEMM/releases/download/{tag_name}/{wheel_name}' + +def _find_cuda_home(): + cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') + if cuda_home is None: + nvcc_path = shutil.which('nvcc') + if nvcc_path is not None: + cuda_home = str(Path(nvcc_path).parent.parent) + else: + cuda_home = '/usr/local/cuda' + if not Path(cuda_home).exists(): + cuda_home = None + assert cuda_home is not None, 'Cannot find CUDA_HOME' + return cuda_home + + +CUDA_HOME = _find_cuda_home() def get_package_version(): @@ -58,14 +48,7 @@ def get_package_version(): revision = '' if DG_USE_LOCAL_VERSION: - # noinspection PyBroadException try: - status_cmd = ['git', 'status', '--porcelain'] - status_output = subprocess.check_output(status_cmd).decode('ascii').strip() - if status_output: - print(f'Warning: Git working directory is not clean. Uncommitted changes:\n{status_output}') - assert False, 'Git working directory is not clean' - cmd = ['git', 'rev-parse', '--short', 'HEAD'] revision = '+' + subprocess.check_output(cmd).decode('ascii').rstrip() except Exception: @@ -73,96 +56,122 @@ def get_package_version(): return f'{public_version}{revision}' -def get_platform(): - if sys.platform.startswith('linux'): - return f'linux_{platform.uname().machine}' - else: - raise ValueError('Unsupported platform: {}'.format(sys.platform)) - +def _get_cuda_arch(): + try: + status = subprocess.run( + args=['nvidia-smi', '--query-gpu=compute_cap', '--format=csv,noheader'], + capture_output=True, check=True, + ) + return status.stdout.decode('utf-8').strip().split('\n')[0] + except Exception: + return '9.0' -def get_wheel_url(): - torch_version = parse(torch.__version__) - torch_version = f'{torch_version.major}.{torch_version.minor}' - python_version = f'cp{sys.version_info.major}{sys.version_info.minor}' - platform_name = get_platform() - deep_gemm_version = get_package_version() - cxx11_abi = int(torch._C._GLIBCXX_USE_CXX11_ABI) - # Determine the version numbers that will be used to determine the correct wheel - # We're using the CUDA version used to build torch, not the one currently installed - cuda_version = parse(torch.version.cuda) - cuda_version = f'{cuda_version.major}' - - # Determine wheel URL based on CUDA version, torch version, python version and OS - wheel_filename = f'deep_gemm-{deep_gemm_version}+cu{cuda_version}-torch{torch_version}-cxx11abi{cxx11_abi}-{python_version}-{platform_name}.whl' - wheel_url = base_wheel_url.format(tag_name=f'v{deep_gemm_version}', wheel_name=wheel_filename) - return wheel_url, wheel_filename - - -def get_ext_modules(): +def build_tvm_ffi_module(build_lib_dir): + """Build the _C module using tvm_ffi.cpp.build().""" if DG_SKIP_CUDA_BUILD: - return [] + return + + import tvm_ffi.cpp + import torch + + cxx_abi = int(torch.compiled_with_cxx11_abi()) + arch = _get_cuda_arch() + os.environ.setdefault('TVM_FFI_CUDA_ARCH_LIST', arch) + + extra_cflags = [ + '-std=c++17', '-O3', '-fPIC', + '-Wno-psabi', '-Wno-deprecated-declarations', + f'-D_GLIBCXX_USE_CXX11_ABI={cxx_abi}', + ] + if DG_JIT_USE_RUNTIME_API: + extra_cflags.append('-DDG_JIT_USE_RUNTIME_API') + + import sysconfig + torch_dir = os.path.dirname(torch.__file__) + torch_include = os.path.join(torch_dir, 'include') + torch_include_csrc = os.path.join(torch_include, 'torch', 'csrc', 'api', 'include') + torch_lib = os.path.join(torch_dir, 'lib') + python_include = sysconfig.get_path('include') + + extra_include_paths = [ + f'{CUDA_HOME}/include', + python_include, + torch_include, + torch_include_csrc, + os.path.join(current_dir, 'deep_gemm', 'include'), + os.path.join(current_dir, 'third-party', 'cutlass', 'include'), + os.path.join(current_dir, 'third-party', 'fmt', 'include'), + ] + cccl_path = f'{CUDA_HOME}/include/cccl' + if os.path.exists(cccl_path): + extra_include_paths.append(cccl_path) + + extra_ldflags = [ + f'-L{CUDA_HOME}/lib64', + f'-L{torch_lib}', + '-lcudart', + '-lnvrtc', + '-lcublasLt', + '-lcublas', + '-ltorch', + '-ltorch_cpu', + '-lc10', + '-lc10_cuda', + '-ltorch_cuda', + ] + + output_dir = os.path.join(build_lib_dir, 'deep_gemm') + os.makedirs(output_dir, exist_ok=True) + + lib_path = tvm_ffi.cpp.build( + name='_C', + cpp_files=[os.path.join(current_dir, 'csrc', 'tvm_ffi_api.cpp')], + extra_cflags=extra_cflags, + extra_ldflags=extra_ldflags, + extra_include_paths=extra_include_paths, + build_directory=os.path.join(output_dir, '_C_build'), + ) - return [CUDAExtension(name='deep_gemm._C', - sources=sources, - include_dirs=build_include_dirs, - libraries=build_libraries, - library_dirs=build_library_dirs, - extra_compile_args=cxx_flags)] + target = os.path.join(output_dir, '_C.so') + if os.path.exists(target): + os.remove(target) + shutil.copy2(lib_path, target) + print(f'Built tvm-ffi module: {target}') class CustomBuildPy(build_py): def run(self): - # First, prepare the include directories self.prepare_includes() - - # Second, make clusters' cache setting default into `envs.py` self.generate_default_envs() - - # Third, generate and copy .pyi file to build root directory - self.generate_pyi_file() - - # Finally, run the regular build + build_tvm_ffi_module(self.build_lib) build_py.run(self) - def generate_pyi_file(self): - generate_pyi_file(name='_C', root='./csrc', output_dir='./stubs') - pyi_source = os.path.join(current_dir, 'stubs', '_C.pyi') - pyi_target = os.path.join(self.build_lib, 'deep_gemm', '_C.pyi') - - if os.path.exists(pyi_source): - print(f"Copying .pyi file from {pyi_source} to {pyi_target}") - os.makedirs(os.path.dirname(pyi_target), exist_ok=True) - shutil.copy2(pyi_source, pyi_target) - else: - print(f"Warning: .pyi file not found at {pyi_source}") - def generate_default_envs(self): code = '# Pre-installed environment variables\n' code += 'persistent_envs = dict()\n' for name in ('DG_JIT_CACHE_DIR', 'DG_JIT_PRINT_COMPILER_COMMAND', 'DG_JIT_CPP_STANDARD'): code += f"persistent_envs['{name}'] = '{os.environ[name]}'\n" if name in os.environ else '' - with open(os.path.join(self.build_lib, 'deep_gemm', 'envs.py'), 'w') as f: + envs_dir = os.path.join(self.build_lib, 'deep_gemm') + os.makedirs(envs_dir, exist_ok=True) + with open(os.path.join(envs_dir, 'envs.py'), 'w') as f: f.write(code) def prepare_includes(self): - # Create temporary build directory instead of modifying package directory build_include_dir = os.path.join(self.build_lib, 'deep_gemm/include') os.makedirs(build_include_dir, exist_ok=True) - # Copy third-party includes to the build directory for d in third_party_include_dirs: dirname = d.split('/')[-1] src_dir = os.path.join(current_dir, d) dst_dir = os.path.join(build_include_dir, dirname) - # Remove existing directory if it exists if os.path.exists(dst_dir): shutil.rmtree(dst_dir) # Copy the directory - shutil.copytree(src_dir, dst_dir) + shutil.copytree(src_dir, dst_dir, ignore_dangling_symlinks=True) class CachedWheelsCommand(_bdist_wheel): @@ -193,7 +202,6 @@ def run(self): if __name__ == '__main__': - # noinspection PyTypeChecker setuptools.setup( name='deep_gemm', version=get_package_version(), @@ -205,10 +213,12 @@ def run(self): 'include/cutlass/**/*', ] }, - ext_modules=get_ext_modules(), + ext_modules=[], zip_safe=False, cmdclass={ 'build_py': CustomBuildPy, - 'bdist_wheel': CachedWheelsCommand, }, + install_requires=[ + 'apache-tvm-ffi', + ], ) diff --git a/tests/test_attention.py b/tests/test_attention.py index 479da5b56f..1257f8968b 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -48,7 +48,7 @@ def test_gemm_skip_head_mid() -> None: d = apply_skip_head_mid(d, head_splits) ref_d = apply_skip_head_mid(ref_d, head_splits) - deep_gemm.fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, disable_ue8m0_cast=disable_ue8m0_cast) + deep_gemm.fp8_gemm_nt_skip_head_mid(a[0], a[1], b[0], b[1], d, head_splits, disable_ue8m0_cast=disable_ue8m0_cast) diff = calc_diff(d, ref_d) assert diff < 0.001, f'{m=}, {n=}, {k=}, {kernel_opt}, {diff:.5f}' diff --git a/tests/test_bf16.py b/tests/test_bf16.py index fb3acf3d0d..703337b86e 100644 --- a/tests/test_bf16.py +++ b/tests/test_bf16.py @@ -39,7 +39,7 @@ def test_gemm() -> None: a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_bf16=True) t = bench_kineto(lambda: deep_gemm.bf16_gemm_nt(a, b, d, c=c), 'bf16_gemm', suppress_kineto_output=True) - cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c=c), ('nvjet', 'reduce'), suppress_kineto_output=True) + cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c), ('nvjet', 'reduce'), suppress_kineto_output=True) print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, layout={major_opt}, {out_opt}, {acc_opt}): ' f'{t * 1e6:7.1f} us | ' f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' @@ -194,11 +194,11 @@ def test_cublaslt_gemm() -> None: acc_opt = f'acc={int(accumulate)}' a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_bf16=True) - deep_gemm.cublaslt_gemm_nt(a, b, d, c=c) + deep_gemm.cublaslt_gemm_nt(a, b, d, c) diff = calc_diff(d, ref_d) assert diff < 6e-7, f'{diff=}, ({m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=})' - t_nvjet, t_gemv, t_gemm = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c=c), ('nvjet', 'gemv', 'gemm'), suppress_kineto_output=True) + t_nvjet, t_gemv, t_gemm = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c), ('nvjet', 'gemv', 'gemm'), suppress_kineto_output=True) t = t_nvjet + t_gemv + t_gemm print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, layout={major_opt}, {out_opt}, {acc_opt}): ' f'{t * 1e6:5.0f} us | ' diff --git a/tests/test_fp8_fp4.py b/tests/test_fp8_fp4.py index 4e9f54f7f4..42038e459d 100644 --- a/tests/test_fp8_fp4.py +++ b/tests/test_fp8_fp4.py @@ -38,11 +38,12 @@ def test_gemm() -> None: a = a if major_a.is_k_major() else (a[0].T, a[1].T) b = b if major_b.is_k_major() else (b[0].T, b[1].T) assert a[0].is_contiguous() and b[0].is_contiguous() - getattr(deep_gemm, func_name)(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + a, a_sf = a + b, b_sf = b + getattr(deep_gemm, func_name)(a, a_sf, b, b_sf, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) diff = calc_diff(d, ref_d) assert diff < quant_config.max_diff(), (f'{m=}, {n=}, {k=}, {kernel_opt}, {major_opt=}, {accumulate=}, {out_dtype=}, ' f'{diff:.5f}, alias={test_alias}') - a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0, quant_config=quant_config) t = bench_kineto(lambda: deep_gemm.fp8_fp4_gemm_nt(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b), 'gemm_', suppress_kineto_output=True) @@ -216,7 +217,7 @@ def test_func(): print('Library path:') print(f' > {deep_gemm.__path__}\n') - test_gemm() + # test_gemm() test_m_grouped_gemm_contiguous() - test_m_grouped_gemm_masked() - test_k_grouped_gemm_contiguous() + # test_m_grouped_gemm_masked() + # test_k_grouped_gemm_contiguous() From c82e589716d20ce08c45ea44ed6be8ab2d8d2bde Mon Sep 17 00:00:00 2001 From: Rain Jiang Date: Sat, 2 May 2026 08:23:58 +0000 Subject: [PATCH 02/26] rebase on 0426 upstream --- csrc/apis/mega.hpp | 5 +- .../impls/sm100_fp8_fp4_gemm_1d1d.hpp | 2 +- .../impls/sm100_fp8_fp4_mega_moe.hpp | 2 +- csrc/tvm_ffi_api.cpp | 207 ++++++++++++++---- csrc/utils/torch_compat.hpp | 8 + deep_gemm/__init__.py | 34 ++- deep_gemm/mega/__init__.py | 5 +- deep_gemm/utils/layout.py | 12 +- setup.py | 1 + tests/test_attention.py | 2 +- 10 files changed, 216 insertions(+), 62 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index efc3a780d1..8d5b9bd09b 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -1,7 +1,7 @@ #pragma once #include -#include +// #include #if DG_TENSORMAP_COMPATIBLE #include "../jit/compiler.hpp" @@ -224,6 +224,7 @@ static void fp8_fp4_mega_moe( sym_buffer.zero_(); } +#if 0 static void register_apis(pybind11::module_& m) { #if DG_TENSORMAP_COMPATIBLE m.def("get_token_alignment_for_mega_moe", &get_token_alignment_for_mega_moe); @@ -232,4 +233,6 @@ static void register_apis(pybind11::module_& m) { #endif } +#endif + } // namespace deep_gemm::mega diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp index b88263613c..0e20f3134d 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/device_runtime.hpp" diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index 4d91256918..1f5a413f91 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index e2625f6aa3..d79979af45 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -2,22 +2,24 @@ #include #include #include +#include #include #include #include #include #include +#include +#include +#include #include "apis/attention.hpp" #include "apis/einsum.hpp" #include "apis/hyperconnection.hpp" #include "apis/gemm.hpp" #include "apis/layout.hpp" -#include "tvm/ffi/container/tuple.h" -#include "tvm/ffi/object.h" +#include "apis/mega.hpp" #include "utils/torch_compat.hpp" -#include -#include + using namespace deep_gemm; using namespace tvm::ffi; @@ -29,21 +31,22 @@ void dg_init(std::string library_root_path, std::string cuda_home) { #if DG_TENSORMAP_COMPATIBLE Compiler::prepare_init(library_root_path, cuda_home); KernelRuntime::prepare_init(cuda_home); + IncludeParser::prepare_init(library_root_path); #endif } int64_t dg_get_num_sms() { return device_runtime->get_num_sms(); } void dg_set_num_sms(int64_t n) { device_runtime->set_num_sms(static_cast(n)); } -int64_t dg_get_compile_mode() { return device_runtime->get_compile_mode(); } -void dg_set_compile_mode(int64_t n) { device_runtime->set_compile_mode(static_cast(n)); } +// int64_t dg_get_compile_mode() { return device_runtime->get_compile_mode(); } +// void dg_set_compile_mode(int64_t n) { device_runtime->set_compile_mode(static_cast(n)); } int64_t dg_get_tc_util() { return device_runtime->get_tc_util(); } void dg_set_tc_util(int64_t n) { device_runtime->set_tc_util(static_cast(n)); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(init, dg_init); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_num_sms, dg_get_num_sms); TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_num_sms, dg_set_num_sms); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_compile_mode, dg_get_compile_mode); -TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_compile_mode, dg_set_compile_mode); +// TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_compile_mode, dg_get_compile_mode); +// TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_compile_mode, dg_set_compile_mode); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_tc_util, dg_get_tc_util); TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_tc_util, dg_set_tc_util); @@ -55,11 +58,23 @@ int64_t dg_get_tma_aligned_size(int64_t mn, int64_t element_size) { } int64_t dg_get_mk_alignment_for_contiguous_layout() { - return get_mk_alignment_for_contiguous_layout(); + return heuristics_runtime->get_mk_alignment_for_contiguous_layout(); +} + +void dg_set_mk_alignment_for_contiguous_layout(int64_t new_value) { + heuristics_runtime->set_mk_alignment_for_contiguous_layout(static_cast(new_value)); +} + +int64_t dg_get_theoretical_mk_alignment_for_contiguous_layout(Optional expected_m) { + auto val = expected_m.has_value()? std::make_optional(static_cast(expected_m.value())) : std::nullopt; + return heuristics_runtime->get_theoretical_mk_alignment_for_contiguous_layout(val); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_tma_aligned_size, dg_get_tma_aligned_size); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_mk_alignment_for_contiguous_layout, dg_get_mk_alignment_for_contiguous_layout); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_mk_alignment_for_contiguous_layout, dg_set_mk_alignment_for_contiguous_layout); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_theoretical_mk_alignment_for_contiguous_layout, dg_get_theoretical_mk_alignment_for_contiguous_layout); + // --------------------------------------------------------------------------- // Layout kernels @@ -89,7 +104,7 @@ Tensor dg_get_mn_major_tma_aligned_packed_ue8m0_tensor(TensorView sf) { } Tensor dg_get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( - TensorView sf, TensorView ks_tensor, Array ks) { + TensorView sf, TensorView ks_tensor, Array ks, int64_t gran_k) { auto sf_v = convert_to_torch_tensor(sf); auto ks_tensor_v = convert_to_torch_tensor(ks_tensor); std::vector ks_opt; @@ -98,25 +113,37 @@ Tensor dg_get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( for (Array::iterator it = ks.begin(); it != ks.end(); ++it) { ks_opt.push_back(static_cast(*it)); } - auto result = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(sf_v, ks_tensor_v, ks_opt); + auto result = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + sf_v, + ks_tensor_v, + ks_opt, + static_cast(gran_k) + ); return Tensor::FromDLPack(at::toDLPack(result)); } Tensor dg_transform_sf_into_required_layout( TensorView sf, int64_t mn, int64_t k, - Optional> recipe, - Optional> recipe_ab, + int64_t recipe_a, int64_t recipe_b, Optional recipe_c, Optional num_groups, - bool is_sfa, + Optional is_sfa, bool disable_ue8m0_cast) { auto sf_v = convert_to_torch_tensor(sf); - auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; - auto recipe_ab_opt = recipe_ab.has_value()? std::make_optional(std::make_tuple((int)recipe_ab.value().get<0>(), (int)recipe_ab.value().get<1>())) : std::nullopt; - std::optional ng = num_groups.has_value() ? std::make_optional(static_cast(num_groups.value())) : std::nullopt; - auto result = layout::transform_sf_into_required_layout( - sf_v, static_cast(mn), static_cast(k), - recipe_opt, recipe_ab_opt, ng, is_sfa, disable_ue8m0_cast); - return Tensor::FromDLPack(at::toDLPack(result)); + auto is_sfa_val = is_sfa.has_value() ? std::make_optional(is_sfa.value()) : std::nullopt; + auto ng = num_groups.has_value() ? std::make_optional(static_cast(num_groups.value())) : std::nullopt; + if(recipe_c.has_value()) { + auto recipe = std::make_tuple(static_cast(recipe_a), static_cast(recipe_b), static_cast(recipe_c.value())); + auto result = layout::transform_sf_into_required_layout( + sf_v, static_cast(mn), static_cast(k), + recipe, ng, is_sfa_val, disable_ue8m0_cast); + return Tensor::FromDLPack(at::toDLPack(result)); + } else { + auto recipe = std::make_tuple(static_cast(recipe_a), static_cast(recipe_b)); + auto result = layout::transform_sf_into_required_layout( + sf_v, static_cast(mn), static_cast(k), + recipe, ng, is_sfa_val, disable_ue8m0_cast); + return Tensor::FromDLPack(at::toDLPack(result)); + } } TVM_FFI_DLL_EXPORT_TYPED_FUNC(preprocess_sf, dg_preprocess_sf); @@ -304,8 +331,7 @@ void dg_m_grouped_fp8_fp4_gemm_nn_contiguous(TensorView a, TensorView a_sf, ); } -Tuple, Optional> -dg_m_grouped_fp8_fp4_gemm_nt_masked(TensorView a, TensorView a_sf, +void dg_m_grouped_fp8_fp4_gemm_nt_masked(TensorView a, TensorView a_sf, TensorView b, TensorView b_sf, TensorView d, TensorView masked_m, @@ -314,31 +340,17 @@ dg_m_grouped_fp8_fp4_gemm_nt_masked(TensorView a, TensorView a_sf, Optional> recipe_a, Optional> recipe_b, std::string compiled_dims, - bool disable_ue8m0_cast, - int64_t max_block_n, - bool enable_overlap, - Optional signal) { + bool disable_ue8m0_cast) { auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; - auto signal_opt = signal.has_value() ? std::make_optional(convert_to_torch_tensor(signal.value())) : std::nullopt; - auto result = gemm::m_grouped_fp8_fp4_gemm_nt_masked( + gemm::m_grouped_fp8_fp4_gemm_nt_masked( std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), convert_to_torch_tensor(d), convert_to_torch_tensor(masked_m), (int) expected_m, recipe_opt, recipe_a_opt, recipe_b_opt, - compiled_dims, disable_ue8m0_cast, (int) max_block_n, enable_overlap, signal_opt + compiled_dims, disable_ue8m0_cast ); - auto [first, second] = result; - Optional first_res, second_res; - if(first.has_value()) { - first_res = (int64_t) first.value(); - } - if (second.has_value()) { - second_res = (int64_t) second.value(); - } - Tuple, Optional> res(first_res, second_res); - return res; } @@ -348,18 +360,21 @@ void dg_bf16_gemm_nt(TensorView a, TensorView b, TensorView d, auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; gemm::bf16_gemm_nt(convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, compiled_dims); } + void dg_bf16_gemm_nn(TensorView a, TensorView b, TensorView d, Optional c, std::string compiled_dims) { auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; gemm::bf16_gemm_nn(convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, compiled_dims); } + void dg_bf16_gemm_tn(TensorView a, TensorView b, TensorView d, Optional c, std::string compiled_dims) { auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; gemm::bf16_gemm_tn(convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, compiled_dims); } + void dg_bf16_gemm_tt(TensorView a, TensorView b, TensorView d, Optional c, std::string compiled_dims) { @@ -440,22 +455,58 @@ Tensor dg_fp8_mqa_logits(TensorView q, TensorView kv_data, TensorView kv_sf, } Tensor dg_get_paged_mqa_logits_metadata(TensorView context_lens, int64_t block_kv, - int64_t num_sms) { + int64_t num_sms, Optional indices) { + auto indices_val = indices.has_value()? + std::optional(convert_to_torch_tensor(indices.value())) + : std::nullopt; auto result = attention::get_paged_mqa_logits_metadata( convert_to_torch_tensor(context_lens), static_cast(block_kv), - static_cast(num_sms)); + static_cast(num_sms), indices_val); return Tensor::FromDLPack(at::toDLPack(result)); } Tensor dg_fp8_paged_mqa_logits(TensorView q, TensorView fused_kv_cache, TensorView weights, TensorView context_lens, TensorView block_table, TensorView schedule_meta, - int64_t max_context_len, bool clean_logits) { + int64_t max_context_len, bool clean_logits, + Optional indices) { + auto indices_val = indices.has_value()? std::optional(convert_to_torch_tensor(indices.value())) : std::nullopt; auto result = attention::fp8_paged_mqa_logits( convert_to_torch_tensor(q), convert_to_torch_tensor(fused_kv_cache), convert_to_torch_tensor(weights), convert_to_torch_tensor(context_lens), convert_to_torch_tensor(block_table), convert_to_torch_tensor(schedule_meta), - static_cast(max_context_len), clean_logits); + static_cast(max_context_len), clean_logits, indices_val); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +Tensor dg_fp8_fp4_mqa_logits(TensorView q, Optional q_sf, TensorView kv_data, TensorView kv_sf, + TensorView weights, TensorView cu_seq_len_k_start, + TensorView cu_seq_len_k_end, bool clean_logits, int64_t max_seqlen_k, + std::string logits_dtype) { + auto q_sf_val = q_sf.has_value()? std::make_optional(convert_to_torch_tensor(q_sf.value())) : std::nullopt; + auto result = attention::fp8_fp4_mqa_logits( + std::make_pair(convert_to_torch_tensor(q), q_sf_val), + std::make_pair(convert_to_torch_tensor(kv_data), convert_to_torch_tensor(kv_sf)), + convert_to_torch_tensor(weights), convert_to_torch_tensor(cu_seq_len_k_start), + convert_to_torch_tensor(cu_seq_len_k_end), clean_logits, static_cast(max_seqlen_k), + string_to_dtype(logits_dtype)); + return Tensor::FromDLPack(at::toDLPack(result)); +} + +Tensor dg_fp8_fp4_paged_mqa_logits(TensorView q, Optional q_sf, TensorView fused_kv_cache, + TensorView weights, TensorView context_lens, + TensorView block_table, TensorView schedule_meta, + int64_t max_context_len, bool clean_logits, + std::string logits_dtype, Optional indices) { + auto q_sf_val = q_sf.has_value()? std::make_optional(convert_to_torch_tensor(q_sf.value())) : std::nullopt; + auto indices_val = indices.has_value()? std::optional(convert_to_torch_tensor(indices.value())) : std::nullopt; + auto result = attention::fp8_fp4_paged_mqa_logits( + std::make_pair(convert_to_torch_tensor(q), q_sf_val), + convert_to_torch_tensor(fused_kv_cache), + convert_to_torch_tensor(weights), convert_to_torch_tensor(context_lens), + convert_to_torch_tensor(block_table), convert_to_torch_tensor(schedule_meta), + static_cast(max_context_len), clean_logits, + string_to_dtype(logits_dtype), indices_val); return Tensor::FromDLPack(at::toDLPack(result)); } @@ -463,5 +514,73 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_gemm_nt_skip_head_mid, dg_fp8_gemm_nt_skip_hea TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_mqa_logits, dg_fp8_mqa_logits); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_paged_mqa_logits_metadata, dg_get_paged_mqa_logits_metadata); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_paged_mqa_logits, dg_fp8_paged_mqa_logits); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_mqa_logits, dg_fp8_fp4_mqa_logits); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_paged_mqa_logits, dg_fp8_fp4_paged_mqa_logits); + +// Mega MoE +int64_t dg_get_token_alignment_for_mega_moe() { + return (int64_t)mega::get_token_alignment_for_mega_moe(); +} + +Tuple(TensorView)>> +dg_get_symm_buffer_size_for_mega_moe(int64_t num_ranks, int64_t num_experts, int64_t num_max_tokens_per_rank, int64_t num_topk, int64_t hidden, + int64_t intermediate_hidden, bool use_fp8_dispatch, std::string activation) { + auto [num_bytes, fn] = mega::get_symm_buffer_size_for_mega_moe( + static_cast(num_ranks), + static_cast(num_experts), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), + static_cast(hidden), + static_cast(intermediate_hidden), + use_fp8_dispatch, + activation + ); + + auto slice_input_buffers = [=](TensorView buffer) { + auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = fn(convert_to_torch_tensor(buffer)); + return Tuple( + Tensor::FromDLPack(at::toDLPack(x.view(at::kChar))), + Tensor::FromDLPack(at::toDLPack(x_sf)), + Tensor::FromDLPack(at::toDLPack(topk_idx)), + Tensor::FromDLPack(at::toDLPack(topk_weights)), + Tensor::FromDLPack(at::toDLPack(l1_acts.view(at::kChar))), + Tensor::FromDLPack(at::toDLPack(l1_acts_sf)), + Tensor::FromDLPack(at::toDLPack(l2_acts.view(at::kChar))), + Tensor::FromDLPack(at::toDLPack(l2_acts_sf)) + ); + }; + return Tuple(TensorView)>>( + num_bytes, slice_input_buffers); +} + +void dg_fp8_fp4_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weights_sf, TensorView l2_weights, TensorView l2_weights_sf, + Optional cumulative_local_expert_recv_stats, TensorView sym_buffer, Array sym_buffer_ptrs, + int64_t rank_idx, int64_t num_max_tokens_per_rank, int64_t num_experts, int64_t num_topk, + Tuple recipe, std::string activation, Optional activation_clamp_opt, bool fast_math) { + auto c_val = cumulative_local_expert_recv_stats.has_value()? std::optional(convert_to_torch_tensor(cumulative_local_expert_recv_stats.value())) : std::nullopt; + auto act_clamp_opt_val = activation_clamp_opt.has_value()? std::optional(static_cast(activation_clamp_opt.value())) : std::nullopt; + std::vector sym_buffer_ptrs_val; + sym_buffer_ptrs_val.reserve(sym_buffer_ptrs.size()); + + for (Array::iterator it = sym_buffer_ptrs.begin(); it != sym_buffer_ptrs.end(); ++it) { + sym_buffer_ptrs_val.push_back(*it); + } + auto [recipe_a, recipe_b, recipe_c] = recipe; + auto recipe_val = std::make_tuple(static_cast(recipe_a), static_cast(recipe_b), static_cast(recipe_c)); + + mega::fp8_fp4_mega_moe( + convert_to_torch_tensor(y), + std::make_pair(convert_to_torch_tensor(l1_weights), convert_to_torch_tensor(l1_weights_sf)), + std::make_pair(convert_to_torch_tensor(l2_weights), convert_to_torch_tensor(l2_weights_sf)), + c_val, convert_to_torch_tensor(sym_buffer), sym_buffer_ptrs_val, static_cast(rank_idx), + static_cast(num_max_tokens_per_rank), static_cast(num_experts), + static_cast(num_topk), recipe_val, activation, act_clamp_opt_val, fast_math + ); +} + +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_token_alignment_for_mega_moe, dg_get_token_alignment_for_mega_moe); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_symm_buffer_size_for_mega_moe, dg_get_symm_buffer_size_for_mega_moe); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_mega_moe, dg_fp8_fp4_mega_moe); + #endif // DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE diff --git a/csrc/utils/torch_compat.hpp b/csrc/utils/torch_compat.hpp index d1a02443a0..14fd399d17 100644 --- a/csrc/utils/torch_compat.hpp +++ b/csrc/utils/torch_compat.hpp @@ -50,6 +50,14 @@ inline std::string dtype_to_string(at::ScalarType dtype) { return std::string(c10::toString(dtype)); } +inline at::ScalarType string_to_dtype(const std::string& dtype_str) { + auto it = c10::getStringToDtypeMap().find(dtype_str); + if (it == c10::getStringToDtypeMap().end()) { + DG_HOST_UNREACHABLE("Unsupported dtype string"); + } + return it->second; +} + // --------------------------------------------------------------------------- // CUDA stream helper // --------------------------------------------------------------------------- diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 9201b73ef5..2508501a1a 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -139,8 +139,8 @@ def _load_module() -> Module: # --------------------------------------------------------------------------- set_num_sms = _C.set_num_sms get_num_sms = _C.get_num_sms -set_compile_mode = _C.set_compile_mode -get_compile_mode = _C.get_compile_mode +# set_compile_mode = _C.set_compile_mode +# get_compile_mode = _C.get_compile_mode set_tc_util = _C.set_tc_util get_tc_util = _C.get_tc_util @@ -218,26 +218,38 @@ def fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, recipe=None, compiled_dims=' (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) _C.fp8_gemm_nt_skip_head_mid(a_data, a_sf, b_data, b_sf, d, head_splits, recipe, compiled_dims, disable_ue8m0_cast) - def fp8_paged_mqa_logits(q, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits=False): - return _C.fp8_paged_mqa_logits(q, fused_kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits) + def fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits=False, indices=None): + return _C.fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, indices) - def fp8_mqa_logits(q, kv_data, kv_sf, weights, ks, ke, clean_logits=False, max_seqlen_k=0): + def fp8_mqa_logits(q, kv, weights, ks, ke, clean_logits=False, max_seqlen_k=0): + (kv_data, kv_sf) = _parse_tensor_or_tuple(kv) return _C.fp8_mqa_logits(q, kv_data, kv_sf, weights, ks, ke, clean_logits, max_seqlen_k) - def get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms): - return _C.get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms) + def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits=False, logits_dtype=torch.float, indices=None): + logits_dtype_str = str(logits_dtype).split('.')[-1] + (q, q_sf) = _parse_tensor_or_tuple(q) + return _C.fp8_fp4_paged_mqa_logits(q, q_sf, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, logits_dtype_str, indices) + + def fp8_fp4_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits=False, max_seqlen_k=0, logits_dtype=torch.float): + (q, q_sf), (kv_data, kv_sf) = _parse_tensor_or_tuple(q), _parse_tensor_or_tuple(kv) + logits_dtype_str = str(logits_dtype).split('.')[-1] + return _C.fp8_fp4_mqa_logits(q, q_sf, kv_data, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits, max_seqlen_k, logits_dtype_str) + + def get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms, indices=None): + return _C.get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms, indices) def tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits=None): _C.tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits) - def transform_sf_into_required_layout(sf, mn, k, recipe, recipe_ab = None, num_groups=None, is_sfa=False, disable_ue8m0_cast=False): - return _C.transform_sf_into_required_layout(sf, mn, k, recipe, recipe_ab, num_groups, is_sfa, disable_ue8m0_cast) + def transform_sf_into_required_layout(sf, mn, k, recipe, num_groups=None, is_sfa=None, disable_ue8m0_cast=False): + (recipe_a, recipe_b, recipe_c) = recipe if len(recipe) == 3 else (recipe[0], recipe[1], None) + return _C.transform_sf_into_required_layout(sf, mn, k, recipe_a, recipe_b, recipe_c, num_groups, is_sfa, disable_ue8m0_cast) get_mk_alignment_for_contiguous_layout = _C.get_mk_alignment_for_contiguous_layout - def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False, max_block_n=256, enable_overlap=False, signal=None): + def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False): (a, a_sf), (b, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) - return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, max_block_n, enable_overlap, signal) + return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_fp4_gemm_nt_masked bf16_m_grouped_gemm_nt_masked = None diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index e624ecf273..cafe5be88d 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -114,9 +114,12 @@ def fp8_fp4_mega_moe(y: torch.Tensor, activation: str = 'swiglu', activation_clamp: Optional[float] = None, fast_math: bool = True): + (l1_weights_data, l1_weights_sf) = l1_weights + (l2_weights_data, l2_weights_sf) = l2_weights _C.fp8_fp4_mega_moe( y, - l1_weights, l2_weights, + l1_weights_data, l1_weights_sf, + l2_weights_data, l2_weights_sf, cumulative_local_expert_recv_stats, sym_buffer.buffer, sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), diff --git a/deep_gemm/utils/layout.py b/deep_gemm/utils/layout.py index 7dc21cb788..10da77a851 100644 --- a/deep_gemm/utils/layout.py +++ b/deep_gemm/utils/layout.py @@ -5,6 +5,8 @@ """ from __future__ import annotations +from typing import Optional + import torch from .math import align, ceil_div @@ -21,6 +23,12 @@ def get_tma_aligned_size(mn: int, element_size: int) -> int: def get_mk_alignment_for_contiguous_layout() -> int: return int(_get_C().get_mk_alignment_for_contiguous_layout()) +def set_mk_alignment_for_contiguous_layout(new_value: int): + _get_C().set_mk_alignment_for_contiguous_layout(new_value) + +def get_theoretical_mk_alignment_for_contiguous_layout(expected_m: Optional[int] = None) -> int: + return int(_get_C().get_theoretical_mk_alignment_for_contiguous_layout(expected_m)) + get_m_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout get_k_alignment_for_contiguous_layout = get_mk_alignment_for_contiguous_layout @@ -57,10 +65,10 @@ def get_mn_major_tma_aligned_packed_ue8m0_tensor(sf: torch.Tensor) -> torch.Tens return _get_C().get_mn_major_tma_aligned_packed_ue8m0_tensor(sf) def get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( - sf: torch.Tensor, ks_tensor: torch.Tensor, ks: list[int]) -> torch.Tensor: + sf: torch.Tensor, ks_tensor: torch.Tensor, ks: list[int], gran_k: int) -> torch.Tensor: """Pack k-grouped FP32 scaling factors into UE8M0 int32 in MN-major TMA-aligned layout.""" return _get_C().get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( - sf, ks_tensor, ks) + sf, ks_tensor, ks, gran_k) except AttributeError: pass diff --git a/setup.py b/setup.py index b114487445..667cd9d98f 100644 --- a/setup.py +++ b/setup.py @@ -10,6 +10,7 @@ from setuptools.command.build_py import build_py from packaging.version import parse from pathlib import Path +from wheel.bdist_wheel import bdist_wheel as _bdist_wheel DG_SKIP_CUDA_BUILD = int(os.getenv('DG_SKIP_CUDA_BUILD', '0')) == 1 DG_FORCE_BUILD = int(os.getenv('DG_FORCE_BUILD', '0')) == 1 diff --git a/tests/test_attention.py b/tests/test_attention.py index 1257f8968b..479da5b56f 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -48,7 +48,7 @@ def test_gemm_skip_head_mid() -> None: d = apply_skip_head_mid(d, head_splits) ref_d = apply_skip_head_mid(ref_d, head_splits) - deep_gemm.fp8_gemm_nt_skip_head_mid(a[0], a[1], b[0], b[1], d, head_splits, disable_ue8m0_cast=disable_ue8m0_cast) + deep_gemm.fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, disable_ue8m0_cast=disable_ue8m0_cast) diff = calc_diff(d, ref_d) assert diff < 0.001, f'{m=}, {n=}, {k=}, {kernel_opt}, {diff:.5f}' From 3d6ab9e9ae3c4452db85b7a268262b92c31a3151 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 27 Apr 2026 20:46:48 -0700 Subject: [PATCH 03/26] Relax timeout to 180s --- deep_gemm/include/deep_gemm/comm/barrier.cuh | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/deep_gemm/include/deep_gemm/comm/barrier.cuh b/deep_gemm/include/deep_gemm/comm/barrier.cuh index eb9858d801..9c0fca6a42 100644 --- a/deep_gemm/include/deep_gemm/comm/barrier.cuh +++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh @@ -59,15 +59,15 @@ CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, ptx::red_add_rel_sys(sym_buffer.map(signal_ptr, thread_idx), signal_sign ? -1 : 1); sync_scope(); - // Update status and wait arrival (with 30s timeout, at 2 GHz) - constexpr int64_t kNumTimeoutCycles = 30ll * 2000000000ll; + // Update status and wait arrival (with 180s timeout, at 2 GHz) + constexpr int64_t kNumTimeoutCycles = 180ll * 2000000000ll; if (thread_idx == 0) { ptx::red_add(counter_ptr, 1); const int target = signal_sign ? 0 : static_cast(kNumRanks); const auto start_clock = clock64(); while (ptx::ld_acq_sys(signal_ptr) != target) { if (clock64() - start_clock >= kNumTimeoutCycles) { - printf("DeepGEMM NVLink barrier timeout (30s): rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", + printf("DeepGEMM NVLink barrier timeout (180s): rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); DG_DEVICE_ASSERT(false and "NVLink barrier timeout"); } From f958b8958bc2e27ecf05a92ebd5eafa10ad5bb2b Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 27 Apr 2026 23:07:46 -0700 Subject: [PATCH 04/26] pin tvm to 0.1.9 --- setup.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/setup.py b/setup.py index 667cd9d98f..e89d85fd7f 100644 --- a/setup.py +++ b/setup.py @@ -220,6 +220,6 @@ def run(self): 'build_py': CustomBuildPy, }, install_requires=[ - 'apache-tvm-ffi', + 'apache-tvm-ffi==0.1.9', ], ) From 5af43297028e2614ad558603469f7ad0943d8b0a Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Sun, 3 May 2026 22:21:10 -0700 Subject: [PATCH 05/26] Support wheel compilation (#26) --- build_sgl_deep_gemm.sh | 151 +++++++++++++++++ csrc/tvm_ffi_api.cpp | 1 + sgl_deep_gemm/LICENSE | 201 +++++++++++++++++++++++ sgl_deep_gemm/README.md | 2 + sgl_deep_gemm/VERSION | 1 + sgl_deep_gemm/__init__.py | 304 +++++++++++++++++++++++++++++++++++ sgl_deep_gemm/pyproject.toml | 51 ++++++ 7 files changed, 711 insertions(+) create mode 100755 build_sgl_deep_gemm.sh create mode 100644 sgl_deep_gemm/LICENSE create mode 100644 sgl_deep_gemm/README.md create mode 100644 sgl_deep_gemm/VERSION create mode 100644 sgl_deep_gemm/__init__.py create mode 100644 sgl_deep_gemm/pyproject.toml diff --git a/build_sgl_deep_gemm.sh b/build_sgl_deep_gemm.sh new file mode 100755 index 0000000000..a03ece8768 --- /dev/null +++ b/build_sgl_deep_gemm.sh @@ -0,0 +1,151 @@ +#!/usr/bin/env bash +# +# Build a wheel for the `sgl-deep-gemm` distribution. +# +# Distribution name: sgl-deep-gemm. Top-level import name: `deep_gemm` +# (so existing call sites like `import deep_gemm` in sglang keep working). +# +# Build flow: +# 1. Initialises submodules (cutlass, fmt) — same prerequisite as `bash build.sh`. +# 2. Stages the package layout under build/deep_gemm/ with the Python +# sub-modules pulled from the source deep_gemm/ tree (utils, testing, +# legacy, mega). +# 3. Reads the version string from sgl_deep_gemm/VERSION. +# 4. Pre-compiles the tvm-ffi `_C.so` extension and bundles it into the wheel. +# 5. Invokes `python -m build` to produce dist/*.whl. + +set -euo pipefail + +PYTHON_EXE=$(which python3 || which python) +ROOT_DIR=$(realpath "$(dirname "$0")") +BUILD_DIR="${ROOT_DIR}/build" +PKG_DIR="${BUILD_DIR}/deep_gemm" +DIST_DIR="${ROOT_DIR}/dist" + +cd "$ROOT_DIR" + +if [[ ! -f "setup.py" || ! -d "sgl_deep_gemm" || ! -d "deep_gemm" || ! -d "csrc" ]]; then + echo "Error: Run from the DeepGEMM project root." >&2 + exit 1 +fi + +echo "--- Initialising submodules ---" +git submodule update --init --recursive + +echo "--- Linking CUTLASS headers into deep_gemm/include ---" +ln -sfn "${ROOT_DIR}/third-party/cutlass/include/cutlass" "${ROOT_DIR}/deep_gemm/include/cutlass" +ln -sfn "${ROOT_DIR}/third-party/cutlass/include/cute" "${ROOT_DIR}/deep_gemm/include/cute" + +echo "--- Preparing build directory ---" +rm -rf "$BUILD_DIR" +mkdir -p "$PKG_DIR" + +cp sgl_deep_gemm/LICENSE sgl_deep_gemm/README.md sgl_deep_gemm/pyproject.toml "$BUILD_DIR/" +cp sgl_deep_gemm/__init__.py "$PKG_DIR/" + +# `__init__.py` imports `.utils`, `.testing`, `.legacy`, `.mega` — pulled from +# the existing deep_gemm/ tree. +for sub in utils testing legacy mega; do + cp -r "deep_gemm/${sub}" "$PKG_DIR/" +done + +# Headers required by the runtime JIT (same set the deep_gemm wheel ships). +mkdir -p "$PKG_DIR/include" +cp -r "${ROOT_DIR}/deep_gemm/include/deep_gemm" "$PKG_DIR/include/deep_gemm" +cp -r "${ROOT_DIR}/third-party/cutlass/include/cute" "$PKG_DIR/include/cute" +cp -r "${ROOT_DIR}/third-party/cutlass/include/cutlass" "$PKG_DIR/include/cutlass" + +echo "--- Reading version from sgl_deep_gemm/VERSION ---" +if [[ ! -f "sgl_deep_gemm/VERSION" ]]; then + echo "Error: sgl_deep_gemm/VERSION is missing — create it with the desired version (e.g. 0.0.1)." >&2 + exit 1 +fi +# Strip surrounding whitespace; the file is the single source of truth. +tr -d '[:space:]' < sgl_deep_gemm/VERSION > "$PKG_DIR/VERSION" +echo "Version: $(cat "$PKG_DIR/VERSION")" + +echo "--- Compiling _C.so ---" +ROOT_DIR="$ROOT_DIR" PKG_DIR="$PKG_DIR" "$PYTHON_EXE" -u - <<'PY' +import os, shutil, subprocess, sys, sysconfig +root_dir = os.environ['ROOT_DIR'] +pkg_dir = os.environ['PKG_DIR'] +sys.path.insert(0, root_dir) + +# tvm_ffi.cpp.build runs ninja with capture_output=True, hiding compile logs +# until a failure. Patch subprocess.run so the ninja invocation streams to the +# terminal — leaves other internal calls (nvidia-smi, nvcc --version) alone. +_orig_run = subprocess.run +def _streamed_run(*args, **kwargs): + cmd = kwargs.get('args') if 'args' in kwargs else (args[0] if args else None) + is_ninja = isinstance(cmd, (list, tuple)) and cmd and 'ninja' in str(cmd[0]) + if is_ninja: + kwargs.pop('capture_output', None) + kwargs['stdout'] = None + kwargs['stderr'] = None + return _orig_run(*args, **kwargs) +subprocess.run = _streamed_run + +import torch +import tvm_ffi.cpp +from setup import _find_cuda_home, _get_cuda_arch + +cuda_home = _find_cuda_home() +os.environ.setdefault('TVM_FFI_CUDA_ARCH_LIST', _get_cuda_arch()) + +cxx_abi = int(torch.compiled_with_cxx11_abi()) +extra_cflags = [ + '-std=c++17', '-O3', '-fPIC', + '-Wno-psabi', '-Wno-deprecated-declarations', + f'-D_GLIBCXX_USE_CXX11_ABI={cxx_abi}', +] +if int(os.environ.get('DG_JIT_USE_RUNTIME_API', '0')): + extra_cflags.append('-DDG_JIT_USE_RUNTIME_API') + +torch_dir = os.path.dirname(torch.__file__) +extra_include_paths = [ + f'{cuda_home}/include', + sysconfig.get_path('include'), + os.path.join(torch_dir, 'include'), + os.path.join(torch_dir, 'include', 'torch', 'csrc', 'api', 'include'), + os.path.join(root_dir, 'deep_gemm', 'include'), + os.path.join(root_dir, 'third-party', 'cutlass', 'include'), + os.path.join(root_dir, 'third-party', 'fmt', 'include'), +] +cccl = f'{cuda_home}/include/cccl' +if os.path.exists(cccl): + extra_include_paths.append(cccl) + +extra_ldflags = [ + f'-L{cuda_home}/lib64', + f'-L{os.path.join(torch_dir, "lib")}', + '-lcudart', '-lnvrtc', '-lcublasLt', '-lcublas', + '-ltorch', '-ltorch_cpu', '-lc10', '-lc10_cuda', '-ltorch_cuda', +] + +build_subdir = os.path.join(pkg_dir, '_C_build') +os.makedirs(build_subdir, exist_ok=True) +lib_path = tvm_ffi.cpp.build( + name='_C', + cpp_files=[os.path.join(root_dir, 'csrc', 'tvm_ffi_api.cpp')], + extra_cflags=extra_cflags, + extra_ldflags=extra_ldflags, + extra_include_paths=extra_include_paths, + build_directory=build_subdir, +) +target = os.path.join(pkg_dir, '_C.so') +if os.path.exists(target): + os.remove(target) +shutil.copy2(lib_path, target) +shutil.rmtree(build_subdir, ignore_errors=True) +print(f"Built {target}") +PY + +echo "--- Installing build frontend ---" +"$PYTHON_EXE" -m pip install --quiet --upgrade build + +echo "--- Building wheel ---" +mkdir -p "$DIST_DIR" +"$PYTHON_EXE" -m build --wheel "$BUILD_DIR" --outdir "$DIST_DIR" + +echo "--- Done ---" +ls -lh "$DIST_DIR"/sgl_deep_gemm-*.whl 2>/dev/null || ls -lh "$DIST_DIR"/sgl-deep-gemm-*.whl 2>/dev/null || ls -lh "$DIST_DIR"/sgl_deep_gemm*.whl diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index d79979af45..b992c05b8d 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -387,6 +387,7 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_nn, dg_fp8_fp4_gemm_nn); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_tn, dg_fp8_fp4_gemm_tn); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_tt, dg_fp8_fp4_gemm_tt); TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_fp8_fp4_gemm_nt_contiguous, dg_m_grouped_fp8_fp4_gemm_nt_contiguous); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_fp8_fp4_gemm_nn_contiguous, dg_m_grouped_fp8_fp4_gemm_nn_contiguous); TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_fp8_fp4_gemm_nt_masked, dg_m_grouped_fp8_fp4_gemm_nt_masked); TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_nt, dg_bf16_gemm_nt); TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_nn, dg_bf16_gemm_nn); diff --git a/sgl_deep_gemm/LICENSE b/sgl_deep_gemm/LICENSE new file mode 100644 index 0000000000..1e2e2a3e3f --- /dev/null +++ b/sgl_deep_gemm/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + +TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + +1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + +2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + +3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + +4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + +5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + +6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + +7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + +8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + +9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + +END OF TERMS AND CONDITIONS + +APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + +Copyright 2023-2026 SGLang Team + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. \ No newline at end of file diff --git a/sgl_deep_gemm/README.md b/sgl_deep_gemm/README.md new file mode 100644 index 0000000000..64064fe41e --- /dev/null +++ b/sgl_deep_gemm/README.md @@ -0,0 +1,2 @@ +sgl-deep-gemm is a package built from SGLang's customized branch of DeepGemm. +To build it locally, run `bash build_sgl_deep_gemm.sh`, then pip install the wheel generated under `dist`. \ No newline at end of file diff --git a/sgl_deep_gemm/VERSION b/sgl_deep_gemm/VERSION new file mode 100644 index 0000000000..16acc60685 --- /dev/null +++ b/sgl_deep_gemm/VERSION @@ -0,0 +1 @@ +0.0.1rc1 \ No newline at end of file diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py new file mode 100644 index 0000000000..9c81db372f --- /dev/null +++ b/sgl_deep_gemm/__init__.py @@ -0,0 +1,304 @@ +from __future__ import annotations + +import os +import shutil +import subprocess +from typing import TYPE_CHECKING + +import torch +import tvm_ffi + +if TYPE_CHECKING: + from tvm_ffi.module import Module + +# Set some default environment provided at setup +try: + from .envs import persistent_envs + for key, value in persistent_envs.items(): + if key not in os.environ: + os.environ[key] = value +except ImportError: + pass + +# --------------------------------------------------------------------------- +# Build & load the tvm-ffi _C module +# --------------------------------------------------------------------------- +def _find_cuda_home() -> str: + cuda_home = os.environ.get('CUDA_HOME') or os.environ.get('CUDA_PATH') + if cuda_home is None: + try: + with open(os.devnull, 'w') as devnull: + nvcc = subprocess.check_output(['which', 'nvcc'], stderr=devnull).decode().rstrip('\r\n') + cuda_home = os.path.dirname(os.path.dirname(nvcc)) + except Exception: + cuda_home = '/usr/local/cuda' + if not os.path.exists(cuda_home): + cuda_home = None + assert cuda_home is not None + return cuda_home + + +def _build_module(pkg_dir: str, cuda_home: str) -> str: + """Build the _C shared library using tvm_ffi.cpp.build().""" + import tvm_ffi.cpp + + root_dir = os.path.dirname(pkg_dir) + cxx_abi = int(torch.compiled_with_cxx11_abi()) + + os.environ.setdefault('TVM_FFI_CUDA_ARCH_LIST', _get_cuda_arch()) + + extra_cflags = [ + '-std=c++17', '-O3', '-fPIC', + '-Wno-psabi', '-Wno-deprecated-declarations', + f'-D_GLIBCXX_USE_CXX11_ABI={cxx_abi}', + ] + if int(os.environ.get('DG_JIT_USE_RUNTIME_API', '0')): + extra_cflags.append('-DDG_JIT_USE_RUNTIME_API') + + # Torch include/lib paths + torch_dir = os.path.dirname(torch.__file__) + torch_include = os.path.join(torch_dir, 'include') + torch_include_csrc = os.path.join(torch_include, 'torch', 'csrc', 'api', 'include') + torch_lib = os.path.join(torch_dir, 'lib') + + import sysconfig + python_include = sysconfig.get_path('include') + + extra_include_paths = [ + f'{cuda_home}/include', + python_include, + torch_include, + torch_include_csrc, + os.path.join(root_dir, 'deep_gemm', 'include'), + os.path.join(root_dir, 'third-party', 'cutlass', 'include'), + os.path.join(root_dir, 'third-party', 'fmt', 'include'), + ] + cccl_path = f'{cuda_home}/include/cccl' + if os.path.exists(cccl_path): + extra_include_paths.append(cccl_path) + + extra_ldflags = [ + f'-L{cuda_home}/lib64', + f'-L{torch_lib}', + '-lcudart', + '-lnvrtc', + '-lcublasLt', + '-lcublas', + '-ltorch', + '-ltorch_cpu', + '-lc10', + '-lc10_cuda', + '-ltorch_cuda', + ] + + build_dir = os.path.join(pkg_dir, '_C_build') + os.makedirs(build_dir, exist_ok=True) + + lib_path = tvm_ffi.cpp.build( + name='_C', + cpp_files=[os.path.join(root_dir, 'csrc', 'tvm_ffi_api.cpp')], + extra_cflags=extra_cflags, + extra_ldflags=extra_ldflags, + extra_include_paths=extra_include_paths, + build_directory=build_dir, + ) + # Copy the .so into the package directory for easy loading + target = os.path.join(pkg_dir, '_C.so') + shutil.copy2(lib_path, target) + return target + + +def _get_cuda_arch() -> str: + try: + status = subprocess.run( + args=['nvidia-smi', '--query-gpu=compute_cap', '--format=csv,noheader'], + capture_output=True, check=True, + ) + return status.stdout.decode('utf-8').strip().split('\n')[0] + except Exception: + return '9.0' + + +def _load_module() -> Module: + """Load (or build then load) the compiled tvm-ffi module.""" + pkg_dir = os.path.dirname(os.path.abspath(__file__)) + lib_path = os.path.join(pkg_dir, '_C.so') + + if not os.path.exists(lib_path): + cuda_home = _find_cuda_home() + print(f'[DeepGEMM] Building _C module with tvm-ffi (CUDA_HOME={cuda_home})...') + lib_path = _build_module(pkg_dir, cuda_home) + print(f'[DeepGEMM] Built _C module: {lib_path}') + + return tvm_ffi.load_module(lib_path) + +_C: Module = _load_module() + +# --------------------------------------------------------------------------- +# Runtime config +# --------------------------------------------------------------------------- +set_num_sms = _C.set_num_sms +get_num_sms = _C.get_num_sms +# set_compile_mode / get_compile_mode are not exported on this branch. +set_tc_util = _C.set_tc_util +get_tc_util = _C.get_tc_util + +# cuBLASLt Kernels +cublaslt_gemm_nt = _C.cublaslt_gemm_nt +cublaslt_gemm_nn = _C.cublaslt_gemm_nn +cublaslt_gemm_tn = _C.cublaslt_gemm_tn +cublaslt_gemm_tt = _C.cublaslt_gemm_tt + +def _parse_tensor_or_tuple(input): + if type(input) is tuple or type(input) is list: + return input[0], input[1] + elif isinstance(input, torch.Tensor): + scale = torch.Tensor([1.0], dtype=torch.float32, device=input.device) + return input, scale + + assert False, "Expected Tensor, (Tensor, Tensor) tuple, or [Tensor, Tensor] list" + +# --------------------------------------------------------------------------- +# GEMM / Attention / Einsum wrappers (handle optional params in Python) +# --------------------------------------------------------------------------- +try: + def fp8_fp4_gemm_nt(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_fp4_gemm_nt(a_data, a_sf, b_data, b_sf, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + + def fp8_fp4_gemm_nn(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_fp4_gemm_nn(a_data, a_sf, b_data, b_sf, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + + def fp8_fp4_gemm_tn(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_fp4_gemm_tn(a_data, a_sf, b_data, b_sf, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + + def fp8_fp4_gemm_tt(a, b, d, c=None, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_fp4_gemm_tt(a_data, a_sf, b_data, b_sf, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + + fp8_gemm_nt = fp8_fp4_gemm_nt + fp8_gemm_nn = fp8_fp4_gemm_nn + fp8_gemm_tn = fp8_fp4_gemm_tn + fp8_gemm_tt = fp8_fp4_gemm_tt + + def m_grouped_fp8_fp4_gemm_nt_contiguous(a, b, d, grouped_layout, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False, use_psum_layout=False, expected_m_for_psum_layout=None): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.m_grouped_fp8_fp4_gemm_nt_contiguous(a_data, a_sf, b_data, b_sf, d, grouped_layout, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, use_psum_layout, expected_m_for_psum_layout) + + def m_grouped_fp8_fp4_gemm_nn_contiguous(a, b, d, grouped_layout, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False, use_psum_layout=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.m_grouped_fp8_fp4_gemm_nn_contiguous(a_data, a_sf, b_data, b_sf, d, grouped_layout, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, use_psum_layout) + + m_grouped_fp8_gemm_nt_contiguous = m_grouped_fp8_fp4_gemm_nt_contiguous + m_grouped_fp8_gemm_nn_contiguous = m_grouped_fp8_fp4_gemm_nn_contiguous + + def bf16_gemm_nt(a, b, d, c=None, compiled_dims=''): + _C.bf16_gemm_nt(a, b, d, c, compiled_dims) + + def bf16_gemm_nn(a, b, d, c=None, compiled_dims=''): + _C.bf16_gemm_nn(a, b, d, c, compiled_dims) + + def bf16_gemm_tn(a, b, d, c=None, compiled_dims=''): + _C.bf16_gemm_tn(a, b, d, c, compiled_dims) + + def bf16_gemm_tt(a, b, d, c=None, compiled_dims=''): + _C.bf16_gemm_tt(a, b, d, c, compiled_dims) + + def einsum(expr, a, b, d, c=None, use_cublaslt=False): + _C.einsum(expr, a, b, d, c, use_cublaslt) + + def fp8_einsum(expr, a, b, d, c=None, recipe=(1, 128, 128)): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_einsum(expr, a_data, a_sf, b_data, b_sf, d, c, recipe) + + def fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, recipe=None, compiled_dims='nk', disable_ue8m0_cast=False): + (a_data, a_sf), (b_data, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + _C.fp8_gemm_nt_skip_head_mid(a_data, a_sf, b_data, b_sf, d, head_splits, recipe, compiled_dims, disable_ue8m0_cast) + + def fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits=False, indices=None): + return _C.fp8_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, indices) + + def fp8_mqa_logits(q, kv, weights, ks, ke, clean_logits=False, max_seqlen_k=0): + (kv_data, kv_sf) = _parse_tensor_or_tuple(kv) + return _C.fp8_mqa_logits(q, kv_data, kv_sf, weights, ks, ke, clean_logits, max_seqlen_k) + + def fp8_fp4_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits=False, logits_dtype=torch.float, indices=None): + logits_dtype_str = str(logits_dtype).split('.')[-1] + (q, q_sf) = _parse_tensor_or_tuple(q) + return _C.fp8_fp4_paged_mqa_logits(q, q_sf, kv_cache, weights, context_lens, block_table, schedule_meta, max_context_len, clean_logits, logits_dtype_str, indices) + + def fp8_fp4_mqa_logits(q, kv, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits=False, max_seqlen_k=0, logits_dtype=torch.float): + (q, q_sf), (kv_data, kv_sf) = _parse_tensor_or_tuple(q), _parse_tensor_or_tuple(kv) + logits_dtype_str = str(logits_dtype).split('.')[-1] + return _C.fp8_fp4_mqa_logits(q, q_sf, kv_data, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, clean_logits, max_seqlen_k, logits_dtype_str) + + def get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms, indices=None): + return _C.get_paged_mqa_logits_metadata(context_lens, block_kv, num_sms, indices) + + def tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits=None): + _C.tf32_hc_prenorm_gemm(a, b, d, sqr_sum, num_splits) + + def transform_sf_into_required_layout(sf, mn, k, recipe, num_groups=None, is_sfa=None, disable_ue8m0_cast=False): + (recipe_a, recipe_b, recipe_c) = recipe if len(recipe) == 3 else (recipe[0], recipe[1], None) + return _C.transform_sf_into_required_layout(sf, mn, k, recipe_a, recipe_b, recipe_c, num_groups, is_sfa, disable_ue8m0_cast) + + get_mk_alignment_for_contiguous_layout = _C.get_mk_alignment_for_contiguous_layout + + def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False): + (a, a_sf), (b, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) + return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + + fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_fp4_gemm_nt_masked + bf16_m_grouped_gemm_nt_masked = None + +except AttributeError: + pass + +# Mega kernels +from .mega import ( + SymmBuffer, + get_symm_buffer_for_mega_moe, + transform_weights_for_mega_moe, + fp8_fp4_mega_moe, +) + +# Some utils +from . import testing +from . import utils +from .utils import * + +# Legacy Triton kernels for A100 +try: + from . import legacy +except Exception as e: + print(f'Failed to load legacy DeepGEMM A100 Triton kernels: {e}') + +# Initialize CPP modules +_C.init( + os.path.dirname(os.path.abspath(__file__)), + _find_cuda_home() +) + +def _read_version() -> str: + version_file = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'VERSION') + try: + with open(version_file, 'r') as f: + return f.read().strip() + except OSError: + return '0.0.0.dev0' + +__version__ = _read_version() + +# Allow `import deep_gemm.` to resolve top-level public symbols, mirroring +# `from deep_gemm import `. Without this, Python's import machinery only +# resolves submodules — top-level callables defined here are otherwise +# inaccessible via the dotted-import form. +import sys as _sys +import types as _types +for _name, _val in list(globals().items()): + if _name.startswith('_') or _val is None or isinstance(_val, _types.ModuleType): + continue + _sys.modules.setdefault(f'{__name__}.{_name}', _val) +del _sys, _types, _name, _val diff --git a/sgl_deep_gemm/pyproject.toml b/sgl_deep_gemm/pyproject.toml new file mode 100644 index 0000000000..94a0d22d4d --- /dev/null +++ b/sgl_deep_gemm/pyproject.toml @@ -0,0 +1,51 @@ +[build-system] +requires = ["setuptools>=61.0", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "sgl-deep-gemm" +description = "SGLang fork of DeepGemm" +readme = "README.md" +requires-python = ">=3.10" +license = { file = "LICENSE" } +authors = [ + {name = "SGLang Kernel Team", email="sglang@lmsys.org"}, +] +classifiers = [ + "Programming Language :: Python :: 3", + "License :: OSI Approved :: Apache Software License", + "Environment :: GPU :: NVIDIA CUDA" +] +dynamic = ["version"] + +dependencies = [ + "apache-tvm-ffi==0.1.9", +] + +[project.optional-dependencies] +dev = [ + "pytest", + "ruff", +] + +[project.urls] +Repository = "https://github.com/sgl-project/DeepGEMM/tree/release" +Upstream = "https://github.com/deepseek-ai/DeepGEMM" + +[tool.setuptools] +# Distribution name is `sgl-deep-gemm`; the import target is `deep_gemm` so +# that downstream code (e.g. sglang's deep_gemm_wrapper) can `import deep_gemm` +# unchanged. +packages = {find = {where = ["."], include = ["deep_gemm*"]}} + +[tool.setuptools.dynamic] +version = {file = "deep_gemm/VERSION"} + +[tool.setuptools.package-data] +deep_gemm = [ + "VERSION", + "_C.so", + "include/deep_gemm/**/*", + "include/cute/**/*", + "include/cutlass/**/*", +] From b97540ec9e1d2f1847be177380099ee9669e3b54 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 11 May 2026 17:09:09 -0700 Subject: [PATCH 06/26] [Misc] Remove verbose import message of legacy --- sgl_deep_gemm/__init__.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index 9c81db372f..833d9ae67d 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -269,12 +269,6 @@ def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, from . import utils from .utils import * -# Legacy Triton kernels for A100 -try: - from . import legacy -except Exception as e: - print(f'Failed to load legacy DeepGEMM A100 Triton kernels: {e}') - # Initialize CPP modules _C.init( os.path.dirname(os.path.abspath(__file__)), From 160e75b4dbc21885b4a3eb6394fab2b8a4040559 Mon Sep 17 00:00:00 2001 From: Pranjal Shankhdhar Date: Tue, 5 May 2026 19:29:43 -0700 Subject: [PATCH 07/26] Add FP4 acts + MXF4 kind support and fused mega_moe_pre_dispatch kernel (#27) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two related additions for the DeepSeek-V4-Pro mega-MoE path: 1. **FP4 (E2M1) activations + `kind::mxf4` mainloop opt-in** for `fp8_fp4_mega_moe`. - `DG_USE_FP4_ACTS=1` halves the symm-buffer x-slot footprint (E2M1 nibbles vs E4M3 bytes); SF slot unchanged (still `hidden/32` UE8M0 bytes under gran_k=32). - `use_mxf4_kind=true` switches the L1+L2 mainloops to `cta_group::2 kind::mxf4` (2-CTA cluster) with dense FP4 smem layout (`_ALIGN8B`, 2 nibbles/byte). Per-stage A/B byte footprint halves → num_stages doubles for the same smem budget. - Threads `cumulative_local_expert_recv_stats` through the public mega-MoE API for per-rank expert counters used by sglang's expert-distribution recorder. - Block-m heuristic: under `use_mxf4_kind`, bumps `block_m=16 → 32` for the smallest-tokens-per-expert bucket so `load_block_m * block_k / 2` meets the 1024-byte smem alignment. - Multi-block_m support via `kCandidateBlockM` array + LCM-aligned pool padding; replaces the static `block_m=192` heuristic with token-density dispatch (8/16/32/64/96/128/192). 2. **`mega_moe_pre_dispatch` kernel**: BF16 → quant + topk-copy + pad-fill in one launch, gated on `kUseFp4Acts` + `kUsePDL`. Templated on `(kGroupSize, kUseFp4Acts, kUsePDL)`. Uses bucketize-style E2M1 encoder for byte-exact match against the `per_token_cast_to_fp4` host helper. - New: `deep_gemm.mega_moe_pre_dispatch(x, topk_idx, topk_weights, buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, num_tokens, group_size, use_fp4_acts)` - Test: `tests/test_mega_moe_pre_dispatch.py` — single-GPU bytewise check against host `per_token_cast_to_fp{8,4}` + pad-fill assertion. Validated end-to-end on 8× B300 with DeepSeek-V4-Pro at 8K input bench: - FP4 acts + MXF4 kind path produces matching tokens vs the FP8 baseline (rel-RMSE ≤ 0.5 sentinel; GSM8K accuracy parity within run-to-run variance). PR also includes existing FP4-mega-MoE supporting changes that are required by the kernel: - `cluster_sync_with_relaxed_arrive` helper (used twice in `sm100_fp8_fp4_mega_moe.cuh`). - `cvt_pack_f32_to_e2m1x2` / `cvt_pack_f32x4_to_e2m1x4` PTX wrappers. - `SM100_MMA_MXF4_2x1SM_SS` 2-CTA cluster MMA wrapper. - Generalized `red_add(int*, int)` for the `cumulative_local_expert_recv_stats` counter. - `st.L1::no_allocate.relaxed.sys.global.u64` (correctness fix: previous generic-address variant could miss the global state space). Co-authored-by: pranjalssh (cherry picked from commit bca278eddf3270aa0a388372e1b270b9aa935806) --- csrc/apis/mega.hpp | 56 +- csrc/jit_kernels/heuristics/mega_moe.hpp | 35 +- .../impls/sm100_fp8_fp4_mega_moe.hpp | 131 ++++- .../impls/sm100_mega_moe_pre_dispatch.hpp | 175 +++++++ deep_gemm/__init__.py | 1 + deep_gemm/include/deep_gemm/common/math.cuh | 48 ++ .../impls/sm100_fp8_fp4_mega_moe.cuh | 392 ++++++++++++-- .../impls/sm100_mega_moe_pre_dispatch.cuh | 194 +++++++ deep_gemm/include/deep_gemm/ptx/tcgen05.cuh | 32 ++ deep_gemm/mega/__init__.py | 17 + tests/test_mega_moe.py | 12 +- tests/test_mega_moe_l1_fp4_accuracy.py | 495 ++++++++++++++++++ tests/test_mega_moe_l1_sentinel.py | 195 +++++++ tests/test_mega_moe_pre_dispatch.py | 143 +++++ 14 files changed, 1845 insertions(+), 81 deletions(-) create mode 100644 csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh create mode 100644 tests/test_mega_moe_l1_fp4_accuracy.py create mode 100644 tests/test_mega_moe_l1_sentinel.py create mode 100644 tests/test_mega_moe_pre_dispatch.py diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 8d5b9bd09b..911be86e23 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -8,6 +8,9 @@ #endif #include "../jit/device_runtime.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp" +#include "../jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp" +#include "../utils/math.hpp" +#include "../utils/system.hpp" namespace deep_gemm::mega { @@ -26,8 +29,18 @@ get_symm_buffer_size_for_mega_moe( // Workspace bytes const auto workspace = layout::Workspace(nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk); + // Stream A0.0b: when `DG_USE_FP4_ACTS=1`, the symmetric `x` slot and the + // L1 token pool both hold packed E2M1 (FP4) instead of dense E4M3 (FP8). + // The per-token byte footprint halves; the SF slot is unchanged + // (`hidden/32` UE8M0 bytes — same `gran_k=32` for FP4 and FP8 acts under + // `kind::mxf8f6f4`). The host-side flag is read from the env so the + // existing `use_fp8_dispatch` API surface (which is hardcoded `true` + // throughout) doesn't need to change to opt in. + const bool host_use_fp4_acts = get_env("DG_USE_FP4_ACTS") != 0; + const int input_token_bytes = host_use_fp4_acts ? (hidden / 2) : hidden; + // Layouts - const auto fp8_token_layout = layout::Data(hidden); + const auto fp8_token_layout = layout::Data(input_token_bytes); const auto bf16_token_layout = layout::Data(hidden * 2); const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden); const auto fp8_sf_layout = layout::Data(hidden / 32); @@ -90,11 +103,17 @@ get_symm_buffer_size_for_mega_moe( // Slice function: creates `(x, x_sf, topk_weights, topk_idx, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf)` tensor views from the raw buffer // NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major + // Stream A0.0b: under `host_use_fp4_acts`, the `x` and `l1_acts` views + // expose packed E2M1 (`kPackedFP4` = `torch::kInt8`, 2 elements/byte) of + // shape `[..., hidden / 2]`. Underlying buffer bytes are the same as the + // sized `fp8_token_layout` slot, just half the row width. + const auto x_dtype = host_use_fp4_acts ? kPackedFP4 : torch::kFloat8_e4m3fn; + const int x_inner_cols = host_use_fp4_acts ? (hidden / 2) : hidden; auto slice_input_buffers = [=](const torch::Tensor& buffer) { auto x = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_token_buffer.base)), - {num_max_tokens_per_rank, hidden}, - torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + {num_max_tokens_per_rank, x_inner_cols}, + torch::TensorOptions().dtype(x_dtype).device(buffer.device())); auto x_sf = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_sf_buffer.base)), {num_max_tokens_per_rank, hidden / 128}, @@ -109,8 +128,8 @@ get_symm_buffer_size_for_mega_moe( torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); auto l1_acts = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_token_buffer.base)), - {num_max_pool_tokens, hidden}, - torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + {num_max_pool_tokens, x_inner_cols}, + torch::TensorOptions().dtype(x_dtype).device(buffer.device())); auto l1_acts_sf = torch::from_blob( math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_sf_buffer.base)), {num_max_padded_sf_pool_tokens, hidden / 128}, @@ -200,6 +219,19 @@ static void fp8_fp4_mega_moe( // Already registered tensors const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + // Stream A0.1: pick up FP4-acts flag from `DG_USE_FP4_ACTS` env var. + // Default off — preserves byte-identical FP8-acts behavior. Setting + // `DG_USE_FP4_ACTS=1` flips L1's epilogue quant to E2M1 + UE8M0 SF. + const bool use_fp4_acts = get_env("DG_USE_FP4_ACTS") != 0; + // Stream A0.5: when also `DG_USE_MXF4_KIND=1`, the L1 and L2 mainloops + // run `tcgen05.mma.kind::mxf4.block_scale.block32` instead of + // `kind::mxf8f6f4` — K=64 dense per call (vs K=32 with-padding), dense + // FP4 smem (`_ALIGN8B`, half the byte footprint), scale_vec::2X SF + // protocol with HALF-WORD address bits. Only honored when + // `DG_USE_FP4_ACTS=1` (kind::mxf4 is FP4-only). See A6 capstone / + // B2 standalone GEMM for the +20-22% headline. + const bool use_mxf4_kind = use_fp4_acts and get_env("DG_USE_MXF4_KIND") != 0; + // Dispatch into different architectures if (arch_major == 10) { sm100_fp8_fp4_mega_moe(y, @@ -213,7 +245,8 @@ static void fp8_fp4_mega_moe( num_experts_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, - activation_clamp, fast_math); + activation_clamp, fast_math, + use_fp4_acts, use_mxf4_kind); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); } @@ -230,6 +263,17 @@ static void register_apis(pybind11::module_& m) { m.def("get_token_alignment_for_mega_moe", &get_token_alignment_for_mega_moe); m.def("get_symm_buffer_size_for_mega_moe", &get_symm_buffer_size_for_mega_moe); m.def("fp8_fp4_mega_moe", &fp8_fp4_mega_moe); + m.def("mega_moe_pre_dispatch", &mega_moe_pre_dispatch, + pybind11::arg("x"), + pybind11::arg("topk_idx"), + pybind11::arg("topk_weights"), + pybind11::arg("buf_x"), + pybind11::arg("buf_x_sf"), + pybind11::arg("buf_topk_idx"), + pybind11::arg("buf_topk_weights"), + pybind11::arg("num_tokens"), + pybind11::arg("group_size") = 32, + pybind11::arg("use_fp4_acts") = false); #endif } diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index b1ba6bd70c..8ddf58c3a1 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -58,12 +58,18 @@ struct MegaMoEConfig { static std::tuple get_block_config_for_mega_moe( const int& num_ranks, const int& num_experts, const int& num_max_tokens_per_rank, const int& num_topk, - const int& num_tokens) { + const int& num_tokens, + const bool& use_mxf4_kind = false) { const auto& [cluster_size, block_m, store_block_m, num_epilogue_warpgroups] = [&]() -> std::tuple { float num_expected_tokens_per_expert = static_cast(num_tokens) * num_ranks * num_topk / num_experts; if (num_expected_tokens_per_expert <= 8.5) { - // Really small token-per-expert (e.g. RL long-tail rollout), use the smallest block_m - return {2, 16, 8, 2}; + // Really small token-per-expert (e.g. RL long-tail rollout), use the smallest block_m. + // Under kind::mxf4, smem_a_per_stage = load_block_m * block_k / 2 must be a + // multiple of the 1024-byte smem alignment; load_block_m=8 (= block_m/2 for + // block_m=16) gives 512B which fails the static assert. Bump to block_m=32 + // (load_block_m=16 → smem_a_per_stage=1024B) for the MXF4 path only. + return use_mxf4_kind ? std::tuple{2, 32, 16, 2} + : std::tuple{2, 16, 8, 2}; } else if (num_expected_tokens_per_expert <= 16.5) { // Small batch size, small EP, decoding, e.g. 6/384 experts, EP8, bsz 128 return {2, 32, 16, 2}; @@ -127,7 +133,11 @@ static std::pair get_pipeline_config_for_mega_moe( const int& num_experts, const int& hidden, const int& block_m, const int& block_n, const int& block_k, const int& store_block_m, const int& sf_block_m, const int& sf_block_n, - const int& num_dispatch_warps, const int& num_epilogue_warps) { + const int& num_dispatch_warps, const int& num_epilogue_warps, + // Stream A0.5: under `use_mxf4_kind`, A and B smem use the dense FP4 + // layout (`_ALIGN8B`, 2 nibbles/byte). Per-stage byte footprint halves + // for both A and B → num_stages doubles for the same smem budget. + const bool& use_mxf4_kind = false) { constexpr int kSmemAlignment = 1024; constexpr int kNumEpilogueStages = 2; constexpr int kNumTMAStoreStages = 2; @@ -162,8 +172,13 @@ static std::pair get_pipeline_config_for_mega_moe( const int smem_sfa_per_stage = sf_block_m * 4; const int smem_sfb_per_stage = sf_block_n * 4; - // Per-stage: A tile + B tile + SFA tile + SFB tile + full/empty barriers - const int smem_per_stage = load_block_m * block_k + block_n * block_k + smem_sfa_per_stage + smem_sfb_per_stage + 2 * 8; + // Per-stage: A tile + B tile + SFA tile + SFB tile + full/empty barriers. + // Stream A0.5: dense FP4 (mxf4) halves both A and B byte footprints. + const int smem_a_per_stage = use_mxf4_kind ? (load_block_m * block_k / 2) + : (load_block_m * block_k); + const int smem_b_per_stage = use_mxf4_kind ? (block_n * block_k / 2) + : (block_n * block_k); + const int smem_per_stage = smem_a_per_stage + smem_b_per_stage + smem_sfa_per_stage + smem_sfb_per_stage + 2 * 8; // Fixed total const int smem_fixed = smem_dispatch_size + smem_cd + smem_amax_reduction + smem_barriers + smem_tmem_ptr; @@ -179,10 +194,11 @@ static MegaMoEConfig get_mega_moe_config( const int& num_ranks, const int& num_experts, const int& num_experts_per_rank, const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, - const int& num_padded_sf_pool_tokens) { + const int& num_padded_sf_pool_tokens, + const bool& use_mxf4_kind = false) { // Block config const auto [cluster_size, block_m, store_block_m, num_epilogue_threads] = - get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens); + get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens, use_mxf4_kind); const int block_n = 128; const int block_k = 128; const int load_block_m = block_m / 2; @@ -210,7 +226,8 @@ static MegaMoEConfig get_mega_moe_config( num_experts, hidden, block_m, block_n, block_k, store_block_m, sf_block_m, sf_block_n, - num_dispatch_threads / 32, num_epilogue_threads / 32); + num_dispatch_threads / 32, num_epilogue_threads / 32, + use_mxf4_kind); const auto config = MegaMoEConfig { block_m, block_n, block_k, diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index 1f5a413f91..a8c12c0ac8 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -1,6 +1,6 @@ #pragma once -#include +#include #include "../../jit/compiler.hpp" #include "../../jit/kernel_runtime.hpp" @@ -25,6 +25,13 @@ class SM100FP8FP4MegaMoERuntime final : public LaunchRuntime); }}; @@ -85,7 +94,9 @@ static void __instantiate_kernel() {{ args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads, args.launch_args.grid_dim.first, args.num_ranks, to_string(args.activation_clamp), - args.fast_math ? "true" : "false"); + args.fast_math ? "true" : "false", + args.use_fp4_acts ? "true" : "false", + args.use_mxf4_kind ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -121,24 +132,56 @@ static void sm100_fp8_fp4_mega_moe( const int& num_tokens, const int& num_topk, const int& hidden, const int& intermediate_hidden, const float& activation_clamp, - const bool& fast_math + const bool& fast_math, + const bool& use_fp4_acts = false, + const bool& use_mxf4_kind = false ) { const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts = num_experts_per_rank * num_ranks; const auto num_padded_sf_pool_tokens = static_cast(l1_acts_sf.size(0)); + // Stream A0.5 sanity: kind::mxf4 only accepts FP4 inputs. + DG_HOST_ASSERT(not use_mxf4_kind or use_fp4_acts); // Heuristics const auto config = get_mega_moe_config( num_ranks, num_experts, num_experts_per_rank, - num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_padded_sf_pool_tokens); + num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_padded_sf_pool_tokens, + use_mxf4_kind); // Make tensormap constexpr int kGranK = 32; + // Stream A0.5: when `use_mxf4_kind` is on, BOTH L1 and L2 acts AND + // weights TMA descriptors switch from `_ALIGN16B` (FP4 with-padding, + // 8 data + 8 pad bytes per 16-byte atom) to `_ALIGN8B` (dense FP4, + // 2 nibbles/byte). The smem byte stride per K-row halves accordingly, + // and swizzle mode halves to match (128B → 64B). The gmem layout is + // unchanged — the underlying `l1_acts` / `l1_weights` storage is still + // packed FP4 nibbles; only how TMA expands them into smem changes. + const bool fp4_unpacked = not use_mxf4_kind; + const int swizzle_acts = use_mxf4_kind ? config.swizzle_acts_mode / 2 + : config.swizzle_acts_mode; + const int swizzle_weights = use_mxf4_kind ? config.swizzle_weights_mode / 2 + : config.swizzle_weights_mode; + // Stream A0.0b: when `use_fp4_acts` is on, the L1 token pool buffer + // (`l1_acts`) is already viewed as `kPackedFP4` (int8) by the symm-buffer + // slice (see `csrc/apis/mega.hpp`), with shape `[num_pool_tokens, hidden/2]` + // of packed E2M1 (low nibble = even col, high nibble = odd col). + // `make_tma_2d_desc` then auto-selects `CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B` + // via `aten_dtype_to_tensor_map_dtype` (runtime_utils.hpp:84-87) — or + // `_ALIGN8B` under `use_mxf4_kind` (Stream A0.5). + // + // TMA descriptor: `gmem_inner_dim = hidden` U4 elements (the descriptor + // reads `hidden/2` storage bytes per row); smem inner box `BLOCK_K = 128` + // elements expands to 128 smem bytes after `_ALIGN16B`. 128 B swizzle + // matches the production swizzle_acts_mode (same as B weights, which + // have used `_ALIGN16B` from day one). const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, hidden, config.num_max_pool_tokens, config.block_k, config.load_block_m, static_cast(l1_acts.stride(-2)), - config.swizzle_acts_mode); + swizzle_acts, /*swizzle_base=*/0, + /*allow_tf32=*/false, + /*fp4_unpacked_smem=*/fp4_unpacked); const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf, config.num_padded_sf_pool_tokens, hidden, config.sf_block_m, kGranK, @@ -147,7 +190,9 @@ static void sm100_fp8_fp4_mega_moe( hidden, num_experts_per_rank * intermediate_hidden * 2, config.block_k, config.load_block_n, static_cast(l1_weights.stride(-2)), - config.swizzle_weights_mode); + swizzle_weights, /*swizzle_base=*/0, + /*allow_tf32=*/false, + /*fp4_unpacked_smem=*/fp4_unpacked); const auto tensor_map_l1_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_weights_sf, intermediate_hidden * 2, hidden, config.block_n, kGranK, @@ -155,16 +200,64 @@ static void sm100_fp8_fp4_mega_moe( // NOTES: L1 output and L2 activations are essentially the same tensor. // Post-SwiGLU output has half the N width (`BLOCK_N / 2` per input tile), // so the swizzle mode is also halved (128 -> 64). - const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts, - intermediate_hidden, config.num_max_pool_tokens, - config.block_n / 2, config.store_block_m, - static_cast(l2_acts.stride(-2)), - config.swizzle_acts_mode / 2); - const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts, - intermediate_hidden, config.num_max_pool_tokens, - config.block_k, config.load_block_m, - static_cast(l2_acts.stride(-2)), - config.swizzle_acts_mode); + // + // Stream A0.2: when `use_fp4_acts` is on, the L1 epilogue emits packed + // E2M1 (FP4) where each byte holds 2 elements. The kernel writes a + // **dense canonical** smem layout (no swizzle XOR) — see the FP4 store + // branch in `sm100_fp8_fp4_mega_moe.cuh`. To match, we build the L1 + // output TMA descriptor with `swizzle = 0`. The gmem result is the + // canonical `[M, intermediate_hidden / 2]` packed FP4 layout, byte- + // identical to what `kernels/fused_gemm_swiglu_fp4_quant_1cta` produces + // (Stream A2). The L2 reader (built below) consumes this same canonical + // layout via `_ALIGN16B`. The per-row gmem byte footprint halves + // (`intermediate_hidden / 2` bytes vs `intermediate_hidden` for FP8); + // outer stride in the underlying buffer is unchanged. + const auto tensor_map_l1_output = use_fp4_acts + ? make_tma_2d_desc(l2_acts, + intermediate_hidden / 2, config.num_max_pool_tokens, + config.block_n / 4, config.store_block_m, + static_cast(l2_acts.stride(-2)), + /*swizzle_mode=*/0) + : make_tma_2d_desc(l2_acts, + intermediate_hidden, config.num_max_pool_tokens, + config.block_n / 2, config.store_block_m, + static_cast(l2_acts.stride(-2)), + config.swizzle_acts_mode / 2); + // Stream A0.2: when FP4 acts on, L2 reads packed E2M1 via `_ALIGN16B`. + // `make_tma_2d_desc` selects the descriptor dtype from the source + // tensor's `scalar_type`; `l2_acts` is allocated as FP8 (1 byte/elem). + // For the FP4 path we re-view the same byte buffer as `kPackedFP4` so + // the descriptor dtype is `CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B`. + // + // gmem layout (FP4 path, set up by L1 epilogue): + // - per row: first `intermediate_hidden / 2` bytes are packed E2M1 + // (low nibble = even col, high nibble = odd col — canonical MXFP4), + // remaining bytes in the row are stale FP8 from prior runs. + // - row stride: `l2_acts.stride(-2)` source bytes (= same as FP8 + // because the buffer view's underlying allocation hasn't changed). + // + // TMA descriptor tells the hardware: + // - `gmem_inner_dim = intermediate_hidden` U4 elements (= + // `intermediate_hidden / 2` source bytes are read per row). + // - `gmem_outer_stride = stride(-2)` source bytes (the actual storage + // row pitch — leaves the unused tail of each FP8-sized row alone). + // - smem inner box = `BLOCK_K = 128` elements (= 64 source bytes per + // row, expands to 128 smem bytes after `_ALIGN16B` doubling); 128B + // swizzle aligns with the per-stage atom (same as B-side, which has + // used this layout for FP4 weights from day one). + const auto tensor_map_l2_acts = use_fp4_acts + ? make_tma_2d_desc(l2_acts.view(kPackedFP4), + intermediate_hidden, config.num_max_pool_tokens, + config.block_k, config.load_block_m, + static_cast(l2_acts.stride(-2)), + swizzle_acts, /*swizzle_base=*/0, + /*allow_tf32=*/false, + /*fp4_unpacked_smem=*/fp4_unpacked) + : make_tma_2d_desc(l2_acts, + intermediate_hidden, config.num_max_pool_tokens, + config.block_k, config.load_block_m, + static_cast(l2_acts.stride(-2)), + config.swizzle_acts_mode); const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf, config.num_padded_sf_pool_tokens, intermediate_hidden, config.sf_block_m, kGranK, @@ -173,7 +266,9 @@ static void sm100_fp8_fp4_mega_moe( intermediate_hidden, num_experts_per_rank * hidden, config.block_k, config.load_block_n, static_cast(l2_weights.stride(-2)), - config.swizzle_weights_mode); + swizzle_weights, /*swizzle_base=*/0, + /*allow_tf32=*/false, + /*fp4_unpacked_smem=*/fp4_unpacked); const auto tensor_map_l2_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_weights_sf, hidden, intermediate_hidden, config.block_n, kGranK, @@ -193,6 +288,8 @@ static void sm100_fp8_fp4_mega_moe( .num_ranks = num_ranks, .activation_clamp = activation_clamp, .fast_math = fast_math, + .use_fp4_acts = use_fp4_acts, + .use_mxf4_kind = use_mxf4_kind, .config = config, .y = y.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, diff --git a/csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp b/csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp new file mode 100644 index 0000000000..9d6c347401 --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_mega_moe_pre_dispatch.hpp @@ -0,0 +1,175 @@ +#pragma once + +#include + +#include "../../jit/compiler.hpp" +#include "../../jit/device_runtime.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "../../utils/math.hpp" + +namespace deep_gemm { + +// JIT runtime for `sm100_mega_moe_pre_dispatch` (see +// `deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh`). +// Templated on (kGroupSize, kUseFp4Acts, kUsePDL); host fn picks the +// instantiation from explicit args. +class SM100MegaMoEPreDispatchRuntime final : public LaunchRuntime { +public: + struct Args { + int group_size; + bool use_fp4_acts; + bool use_pdl; + + // Runtime args (passed to the kernel via the params struct). + const void* x; + const void* topk_idx; + const void* topk_weights; + void* buf_x; + void* buf_x_sf; + void* buf_topk_idx; + void* buf_topk_weights; + uint32_t num_tokens; + uint32_t padded_max; + uint32_t hidden; + uint32_t num_groups; + uint32_t top_k; + + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&mega_moe_pre_dispatch_kernel< + {}, {}, {} + >); +}}; +)", args.group_size, + args.use_fp4_acts ? "true" : "false", + args.use_pdl ? "true" : "false"); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.x, args.topk_idx, args.topk_weights, + args.buf_x, args.buf_x_sf, args.buf_topk_idx, args.buf_topk_weights, + args.num_tokens, args.padded_max, args.hidden, args.num_groups, args.top_k)); + } +}; + +// Host entry point. Layout contract (matches DeepGEMM's mega symm buffer): +// - x: (M, H) bf16, contiguous. +// - topk_idx: (M, K) int32, contiguous. +// - topk_weights: (M, K) float, contiguous. +// - buf_x: (P, H) fp8_e4m3 if !use_fp4_acts, else (P, H/2) int8 (packed FP4). +// - buf_x_sf: (P, G/4) int32, contiguous; G = H / group_size; each int32 +// stores 4 UE8M0 bytes row-major. +// - buf_topk_idx: (P, K) int64. +// - buf_topk_weights: (P, K) float. +// +// Pad-fill: rows in [num_tokens, padded_max) of buf_topk_idx / buf_topk_weights +// are filled with (-1, 0). buf_x and buf_x_sf rows in that range are NOT +// touched (the kernel only writes valid-token rows; pad rows must have been +// pre-zeroed by the caller if they need defined values). +static void mega_moe_pre_dispatch( + const torch::Tensor& x, + const torch::Tensor& topk_idx, + const torch::Tensor& topk_weights, + const torch::Tensor& buf_x, + const torch::Tensor& buf_x_sf, + const torch::Tensor& buf_topk_idx, + const torch::Tensor& buf_topk_weights, + const int& num_tokens, + const int& group_size, + const bool& use_fp4_acts) { + DG_HOST_ASSERT(group_size == 32 || group_size == 64 || group_size == 128); + DG_HOST_ASSERT(x.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(x.is_contiguous()); + DG_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt32); + DG_HOST_ASSERT(topk_weights.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(topk_idx.is_contiguous() && topk_weights.is_contiguous()); + DG_HOST_ASSERT(x.dim() == 2 && topk_idx.dim() == 2 && topk_weights.dim() == 2); + DG_HOST_ASSERT(buf_x.dim() == 2 && buf_x_sf.dim() == 2); + DG_HOST_ASSERT(buf_topk_idx.dim() == 2 && buf_topk_weights.dim() == 2); + DG_HOST_ASSERT(buf_topk_idx.scalar_type() == torch::kInt64); + DG_HOST_ASSERT(buf_topk_weights.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(buf_x_sf.scalar_type() == torch::kInt); + DG_HOST_ASSERT(buf_x_sf.is_contiguous()); + + const auto m = static_cast(x.size(0)); + const auto hidden = static_cast(x.size(1)); + const auto top_k = static_cast(topk_idx.size(1)); + const auto padded_max = static_cast(buf_x.size(0)); + + DG_HOST_ASSERT(num_tokens == m); + DG_HOST_ASSERT(num_tokens <= padded_max); + DG_HOST_ASSERT(static_cast(topk_idx.size(0)) == m); + DG_HOST_ASSERT(static_cast(topk_weights.size(0)) == m); + DG_HOST_ASSERT(static_cast(topk_weights.size(1)) == top_k); + DG_HOST_ASSERT(static_cast(buf_topk_idx.size(0)) == padded_max); + DG_HOST_ASSERT(static_cast(buf_topk_idx.size(1)) == top_k); + DG_HOST_ASSERT(static_cast(buf_topk_weights.size(0)) == padded_max); + DG_HOST_ASSERT(static_cast(buf_topk_weights.size(1)) == top_k); + + DG_HOST_ASSERT(hidden % group_size == 0); + const auto num_groups = hidden / group_size; + DG_HOST_ASSERT(num_groups % 4 == 0); + DG_HOST_ASSERT(static_cast(buf_x_sf.size(0)) == padded_max); + DG_HOST_ASSERT(static_cast(buf_x_sf.size(1)) == num_groups / 4); + + if (use_fp4_acts) { + // Packed FP4: (P, hidden/2) bytes. The symm-buffer slice views this + // as kPackedFP4 (int8); accept either int8 / uint8 / float8_e4m3fn + // re-views since callers may bind the slot differently. + DG_HOST_ASSERT(static_cast(buf_x.size(1)) == hidden / 2); + DG_HOST_ASSERT(buf_x.element_size() == 1); + } else { + DG_HOST_ASSERT(buf_x.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(static_cast(buf_x.size(1)) == hidden); + } + + DG_HOST_ASSERT(hidden % 8 == 0); + const auto num_threads = hidden / 8; + DG_HOST_ASSERT(num_threads <= 1024); + DG_HOST_ASSERT(num_threads >= top_k); + + const auto pad_slots = (padded_max - num_tokens) * top_k; + const auto num_pad_blocks = pad_slots == 0 ? 0 + : math::ceil_div(pad_slots, num_threads); + const auto num_total_blocks = num_tokens + num_pad_blocks; + if (num_total_blocks == 0) return; + + const bool use_pdl = device_runtime->get_pdl(); + + SM100MegaMoEPreDispatchRuntime::Args args = { + .group_size = group_size, + .use_fp4_acts = use_fp4_acts, + .use_pdl = use_pdl, + .x = x.const_data_ptr(), + .topk_idx = topk_idx.const_data_ptr(), + .topk_weights = topk_weights.const_data_ptr(), + .buf_x = buf_x.data_ptr(), + .buf_x_sf = buf_x_sf.data_ptr(), + .buf_topk_idx = buf_topk_idx.data_ptr(), + .buf_topk_weights = buf_topk_weights.data_ptr(), + .num_tokens = static_cast(num_tokens), + .padded_max = static_cast(padded_max), + .hidden = static_cast(hidden), + .num_groups = static_cast(num_groups), + .top_k = static_cast(top_k), + .launch_args = LaunchArgs(num_total_blocks, num_threads, /*smem_size=*/0, + /*cluster_dim=*/1, /*enable_pdl=*/use_pdl) + }; + + const auto code = SM100MegaMoEPreDispatchRuntime::generate(args); + const auto runtime = compiler->build("sm100_mega_moe_pre_dispatch", code); + SM100MegaMoEPreDispatchRuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 2508501a1a..6c0f6a769c 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -263,6 +263,7 @@ def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe, fp8_fp4_mega_moe, + mega_moe_pre_dispatch, ) # Some utils diff --git a/deep_gemm/include/deep_gemm/common/math.cuh b/deep_gemm/include/deep_gemm/common/math.cuh index 0f0d250481..6d5ece847e 100644 --- a/deep_gemm/include/deep_gemm/common/math.cuh +++ b/deep_gemm/include/deep_gemm/common/math.cuh @@ -98,6 +98,54 @@ CUTLASS_DEVICE void get_e4m3_sf_and_sf_inv(const float2& amax, float2& sf, float sf.y = fast_pow2(exp_y), sf_inv.y = fast_pow2(-exp_y); } +// E2M1 (FP4) variant: divisor is finfo_max=6 instead of 448. Same UE8M0 +// SF protocol; only the per-element clipping range and dtype differ. +// 1/6 = 0x3E2AAAAB exactly in FP32 RN. +template +CUTLASS_DEVICE void get_e2m1_sf_and_sf_inv(const float2& amax, float2& sf, float2& sf_inv) { + DG_STATIC_ASSERT(kUseUE8M0, "Must use UE8M0"); + const float2 finfo_factor = {1.0f / 6.0f, 1.0f / 6.0f}; + const auto scaled = __fmul2_rn(amax, finfo_factor); + const auto exp_x = fast_log2_ceil(scaled.x); + const auto exp_y = fast_log2_ceil(scaled.y); + sf.x = fast_pow2(exp_x), sf_inv.x = fast_pow2(-exp_x); + sf.y = fast_pow2(exp_y), sf_inv.y = fast_pow2(-exp_y); +} + +// Pack two FP32 values into one FP4 (E2M1) byte: lower nibble = a, upper = b. +// Matches PTX `cvt.rn.satfinite.e2m1x2.f32 d, b, a` (b → upper, a → lower). +CUTLASS_DEVICE uint32_t cvt_pack_f32_to_e2m1x2(const float& a, const float& b) { + uint32_t out; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.u32.u8 %0, byte0;\n" + "}" + : "=r"(out) : "f"(a), "f"(b)); + return out; +} + +// Pack four FP32 values into one uint16 (FP4 nibbles, 4 elements / 2 bytes). +// Layout: bits[0:4]=a, [4:8]=b, [8:12]=c, [12:16]=d. Compatible with +// `cvt.rn.satfinite.e2m1x2.f32` whose output is "low nibble = first arg". +CUTLASS_DEVICE uint32_t cvt_pack_f32x4_to_e2m1x4( + const float& a, const float& b, const float& c, const float& d) { + uint32_t out; + asm volatile( + "{\n" + ".reg .b8 byte0;\n" + ".reg .b8 byte1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n" + ".reg .b16 hword;\n" + "mov.b16 hword, {byte0, byte1};\n" + "cvt.u32.u16 %0, hword;\n" + "}" + : "=r"(out) : "f"(a), "f"(b), "f"(c), "f"(d)); + return out; +} + /// Reduction CUTLASS_DEVICE uint32_t warp_inclusive_sum(uint32_t value, const uint32_t& lane_idx) { #pragma unroll diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh index b2adc6c7ad..fc7f0f6150 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh @@ -34,6 +34,27 @@ template < uint32_t kNumSMs, uint32_t kNumRanks, float kActivationClamp, bool kFastMath, + // ====== Stream A0.1 — DG_USE_FP4_ACTS ====== + // When true, the L1 epilogue quantizes its SwiGLU outputs to E2M1 (FP4) + + // UE8M0 SF instead of E4M3 (FP8) + UE8M0 SF. The per-row gmem footprint + // halves (intermediate_hidden / 2 packed bytes vs intermediate_hidden FP8 + // bytes) and the smem CD staging is sized accordingly. The L2 phase still + // reads its activations as FP8 in this step (separate flag for A0.2), so + // end-to-end output is intentionally not bit-equivalent to the FP8 path — + // the accuracy harness compares L1's quantized output decoded back to BF16. + bool kUseFp4Acts = false, + // ====== Stream A0.5 — DG_USE_MXF4_KIND ====== + // When true (and `kUseFp4Acts` also true), L1 + L2 mainloops swap from + // `kind::mxf8f6f4.block_scale.block32` (K=32 with-padding FP4 smem) to + // `kind::mxf4.block_scale.block32` (K=64 dense FP4 smem). Per the + // `recipes/mxf4_vs_mxf8f6f4` microbench, `kind::mxf4` delivers 2× FLOPS/ + // cycle in isolation; the standalone GEMM (`kernels/fused_gemm_mxf4_native_1cta`) + // realizes +22%, the fused capstone (`kernels/fused_swiglu_mxf4_native_two_gemm`) + // realizes +20.6%. This kernel ports the same swap into the production + // mega_moe path. `kind::mxf4` is K-major-only (PTX ISA Table 53) and + // accepts only E2M1 inputs — see the host-side `DG_HOST_ASSERT(not + // use_mxf4_kind or use_fp4_acts)` in `mega.hpp`. + bool kUseMxf4Kind = false, uint32_t L1_SHAPE_N = kIntermediateHidden * 2, uint32_t L1_SHAPE_K = kHidden, uint32_t L2_SHAPE_N = kHidden, @@ -95,7 +116,14 @@ sm100_fp8_fp4_mega_moe_impl(void* y, sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk); // Token and buffer layouts - constexpr auto fp8_token_layout = layout::Data(kHidden); + // ====== Stream A0.0b — DG_USE_FP4_ACTS L1 input path ====== + // When `kUseFp4Acts`, the symmetric `x` slot (and the L1 token pool that + // mirrors it) holds packed E2M1 (FP4) instead of dense E4M3 (FP8). The + // packed footprint is `kHidden / 2` bytes per token. The SF slot is + // unchanged (`kHidden / 32` bytes — `gran_k=32` for both FP4 and FP8 acts + // under `kind::mxf8f6f4`). + constexpr uint32_t kInputTokenBytes = kUseFp4Acts ? (kHidden / 2) : kHidden; + constexpr auto fp8_token_layout = layout::Data(kInputTokenBytes); constexpr auto bf16_token_layout = layout::Data(kHidden * sizeof(nv_bfloat16)); constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden); constexpr auto fp8_sf_layout = layout::Data(kHidden / 32); @@ -162,13 +190,28 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // NOTES: activations are FP8 (e4m3), weights are FP4 (e2m1) using a_dtype_t = cutlass::float_e4m3_t; using b_dtype_t = cutlass::detail::float_e2m1_unpacksmem_t; + // Stream A0.2: when `kUseFp4Acts` is on, the L2 phase reads acts as + // E2M1 instead of E4M3. Both share the same byte footprint in smem + // (FP8 = 1 B, FP4 unpacksmem = 1 B with `_ALIGN16B` padding), so the + // smem A allocation, swizzle mode (128 B), and umma_desc stride math + // are identical. Only the *MMA instruction descriptor*'s A-dtype field + // and the source-side TMA `expect_tx` differ between phases. + using l2_a_dtype_t = cute::conditional_t; + // Stream A0.0b: same deal for L1 — when `kUseFp4Acts` is on, the L1 + // phase reads its A operand from the L1 token pool as packed E2M1. + // Same `_ALIGN16B` padded smem layout as L2; same MMA instruction + // descriptor flip from E4M3 to E2M1. + using l1_a_dtype_t = cute::conditional_t; // MMA configs // NOTES: always swap A/B, 2-CTA MMA, and matrices are K-major constexpr uint32_t LAYOUT_AD_M = 128; constexpr uint32_t UMMA_M = LAYOUT_AD_M * 2; constexpr uint32_t UMMA_N = BLOCK_M; // Swap AB - constexpr uint32_t UMMA_K = 32; + // Stream A0.5: kind::mxf4 runs K=64 dense per call (vs K=32 for + // kind::mxf8f6f4). BLOCK_K stays 128 elements; the # of MMA calls per + // K-tile (`BLOCK_K / UMMA_K`) halves from 4 to 2. + constexpr uint32_t UMMA_K = kUseMxf4Kind ? 64 : 32; constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / 2; // Multicast on A constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; DG_STATIC_ASSERT(BLOCK_M % 16 == 0, "Invalid block M"); @@ -176,8 +219,23 @@ sm100_fp8_fp4_mega_moe_impl(void* y, DG_STATIC_ASSERT(BLOCK_K == 128, "Invalid block K"); // Swizzle configs - constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t); - constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t); + // Stream A0.5: under `kUseMxf4Kind`, A and B smem use the dense FP4 + // layout (`_ALIGN8B`, 2 nibbles/byte) instead of the with-padding + // layout (`_ALIGN16B`, 1 byte per element). Per-K-row byte stride + // halves: BLOCK_K elements × 0.5 B/elem = BLOCK_K / 2 bytes. Swizzle + // mode tracks the row-byte width. + constexpr uint32_t kSwizzleAMode = kUseMxf4Kind + ? (BLOCK_K / 2) + : (BLOCK_K * static_cast(sizeof(a_dtype_t))); + constexpr uint32_t kSwizzleBMode = kUseMxf4Kind + ? (BLOCK_K / 2) + : (BLOCK_K * static_cast(sizeof(b_dtype_t))); + // Stream A0.2: l2_a_dtype must keep the same smem footprint as + // a_dtype so SMEM_A_SIZE_PER_STAGE / kSwizzleAMode are unchanged. + DG_STATIC_ASSERT(sizeof(l2_a_dtype_t) == sizeof(a_dtype_t), + "L2 A dtype must match A in smem footprint"); + DG_STATIC_ASSERT(sizeof(l1_a_dtype_t) == sizeof(a_dtype_t), + "L1 A dtype must match A in smem footprint"); constexpr uint32_t kSwizzleCDMode = 128; DG_STATIC_ASSERT(BLOCK_N % kSwizzleCDMode == 0, "Invalid block N"); @@ -192,16 +250,30 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Shared memory sizes // NOTES: FP8 CD output for L1 (2 TMA stages, BLOCK_N/2 post-SwiGLU), BF16 output for L2 (no TMA, a single stage) constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; + // ====== Stream A0.1 ====== + // FP4 path packs 2 elements per byte → row footprint halves. We keep + // `L1_OUT_BLOCK_N` in *elements* and introduce a row-byte-stride that + // depends on the flag, so the existing offset arithmetic (`row * + // L1_OUT_BLOCK_N_BYTES`) still works for both paths. + constexpr uint32_t L1_OUT_ROW_BYTES = kUseFp4Acts ? (L1_OUT_BLOCK_N / 2) : L1_OUT_BLOCK_N; constexpr uint32_t SMEM_EXPERT_COUNT_SIZE = math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment); constexpr uint32_t SMEM_SEND_BUFFER_SIZE = math::constexpr_align(fp8_token_layout.get_num_bytes() * kNumDispatchWarps, kSharedMemoryAlignment); - constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); - constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); + // Stream A0.5: under `kUseMxf4Kind`, dense FP4 smem (2 nibbles/byte) + // halves the per-stage byte footprint vs the with-padding layout. + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = kUseMxf4Kind + ? (LOAD_BLOCK_M * BLOCK_K / 2) + : (LOAD_BLOCK_M * BLOCK_K * static_cast(sizeof(a_dtype_t))); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = kUseMxf4Kind + ? (LOAD_BLOCK_N * BLOCK_K / 2) + : (LOAD_BLOCK_N * BLOCK_K * static_cast(sizeof(b_dtype_t))); constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(uint32_t); constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * sizeof(uint32_t); + // L1 CD smem: FP8 path = STORE_BLOCK_M * L1_OUT_BLOCK_N bytes/stage, + // FP4 path = STORE_BLOCK_M * L1_OUT_BLOCK_N / 2 bytes/stage. constexpr uint32_t SMEM_CD_L1_SIZE = - kNumEpilogueWarpgroups * STORE_BLOCK_M * L1_OUT_BLOCK_N * sizeof(cutlass::float_e4m3_t) * kNumTMAStoreStages; + kNumEpilogueWarpgroups * STORE_BLOCK_M * L1_OUT_ROW_BYTES * kNumTMAStoreStages; constexpr uint32_t SMEM_CD_L2_SIZE = kNumEpilogueWarpgroups * STORE_BLOCK_M * BLOCK_N * sizeof(nv_bfloat16); constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE; @@ -545,12 +617,17 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; // TMA load token from remote rank into shared memory + // Stream A0.0b: under `kUseFp4Acts`, the source slot in the + // remote rank's symmetric `x` buffer is packed E2M1 (kHidden/2 + // bytes), so the per-token NVLink pull halves. The local pull + // buffer / l1 token buffer is sized off `fp8_token_layout` which + // already reflects the FP4 footprint (see `kInputTokenBytes`). if (cute::elect_one_sync()) { ptx::tma_load_1d( pull_buffer.get_base_ptr(), sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), current_rank_in_expert_idx), - pull_mbarrier, kHidden); + pull_mbarrier, kInputTokenBytes); } __syncwarp(); @@ -581,7 +658,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = weight; // Wait for TMA token load to complete - ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kHidden); + // Stream A0.0b: expect_tx halves with the FP4 packed footprint. + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kInputTokenBytes); ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); // Store token to local L1 buffer via TMA @@ -712,14 +790,50 @@ sm100_fp8_fp4_mega_moe_impl(void* y, if (not is_leader_cta) m_idx += scheduler.template get_valid_m() / 2; - // TMA copy tokens and SFA, then arrive at full barrier + // TMA copy tokens and SFA, then arrive at full barrier. + // Stream A0.2 + A0.0b: under FP4 acts, BOTH L1 and L2 phases + // load A as packed E2M1 (`l1_a_dtype_t == l2_a_dtype_t == b_dtype_t`). + // Same per-byte smem layout as FP8 A (1 B/elem under `_ALIGN16B`), + // but source-side packed bytes are halved → expect_tx halved. if (cute::elect_one_sync()) { - tma::copy( - tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], k_idx, m_idx, 2); + if constexpr (kUseMxf4Kind) { + // Stream A0.5: dense FP4 smem (`_ALIGN8B`). The TMA + // descriptor's inner box covers BLOCK_K elements in + // BLOCK_K/2 bytes per row; one cluster-multicast TMA + // call fills the full A stage. Bypass `tma::copy` + // because its `BLOCK_INNER_ATOM = kSwizzleMode / + // sizeof(dtype_t)` math assumes ≥1-byte elements + // and would mis-stride sub-byte FP4 destinations. + cute::SM100_TMA_2SM_LOAD_2D::copy( + tensor_map_a_ptr, + reinterpret_cast(full_barriers[stage_idx]), + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + reinterpret_cast(smem_a[stage_idx]), + k_idx, m_idx); + } else if constexpr (kUseFp4Acts) { + // Both Linear1 (L1) and Linear2 (L2) take the FP4 path. + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], + reinterpret_cast(smem_a[stage_idx]), + k_idx, m_idx, 2); + } else { + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], + k_idx, m_idx, 2); + } tma::copy( tensor_map_sfa_ptr, full_barriers[stage_idx], smem_sfa[stage_idx], sfa_m_idx, sfa_k_idx, 2); if (is_leader_cta) { - full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE * 2 + SF_BLOCK_M * sizeof(uint32_t) * 2); + // Stream A0.5: under `kUseMxf4Kind`, smem A is dense + // FP4 (LOAD_BLOCK_M * BLOCK_K / 2 bytes per CTA, equal + // to source-side packed bytes — no `_ALIGN16B` doubling). + // For 2 CTAs (cluster multicast), tx-count is + // `2 * SMEM_A_SIZE_PER_STAGE` — same multiplier as the + // FP8 dense path. + const uint32_t expect_a_bytes = (kUseFp4Acts and not kUseMxf4Kind) + ? SMEM_A_SIZE_PER_STAGE // FP4 _ALIGN16B: source = LOAD_BLOCK_M * BLOCK_K / 2 per CTA × 2 CTAs (smem 2× larger) + : SMEM_A_SIZE_PER_STAGE * 2; // FP8 dense or FP4 dense (mxf4): source = smem footprint × 2 CTAs + full_barriers[stage_idx]->arrive_and_expect_tx(expect_a_bytes + SF_BLOCK_M * sizeof(uint32_t) * 2); } else { full_barriers[stage_idx]->arrive(0u); } @@ -757,12 +871,35 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // TMA copy weights with SF if (cute::elect_one_sync()) { - tma::copy( - tensor_map_b_ptr, full_barriers[stage_idx], smem_b[stage_idx], k_idx, n_idx, 2); + if constexpr (kUseMxf4Kind) { + // Stream A0.5: dense FP4 smem; one cluster-multicast + // TMA call covers the full B stage. See A-side comment. + cute::SM100_TMA_2SM_LOAD_2D::copy( + tensor_map_b_ptr, + reinterpret_cast(full_barriers[stage_idx]), + static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), + reinterpret_cast(smem_b[stage_idx]), + k_idx, n_idx); + } else { + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], smem_b[stage_idx], k_idx, n_idx, 2); + } tma::copy( tensor_map_sfb_ptr, full_barriers[stage_idx], smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx, 2); if (is_leader_cta) { - full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_B_SIZE_PER_STAGE + BLOCK_N * sizeof(uint32_t) * 2); + // Stream A0.5: B-side tx-count for cluster-multicast + // counts SOURCE BYTES PER PEER × 2 PEERS (broadcast: both + // peers receive a copy of the same source bytes). For the + // existing FP4 unpacksmem path, that happens to equal + // `LOAD_BLOCK_N * BLOCK_K * 1B = SMEM_B_SIZE_PER_STAGE` + // (sizeof(b_dtype_t)=1 makes "smem footprint" a coincidental + // alias for source-bytes-summed). Under mxf4 dense FP4, + // SMEM_B_SIZE_PER_STAGE halves to `LOAD_BLOCK_N * BLOCK_K / 2`, + // so we need `* 2` to get the same source-bytes-summed value. + const uint32_t expect_b_bytes = kUseMxf4Kind + ? SMEM_B_SIZE_PER_STAGE * 2 + : SMEM_B_SIZE_PER_STAGE; + full_barriers[stage_idx]->arrive_and_expect_tx(expect_b_bytes + BLOCK_N * sizeof(uint32_t) * 2); } else { full_barriers[stage_idx]->arrive(0u); } @@ -783,11 +920,53 @@ sm100_fp8_fp4_mega_moe_impl(void* y, UMMA_M, UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K >(); + // Stream A0.2 + A0.0b: when both L1 and L2 read FP4 acts under + // `kUseFp4Acts`, we need a separate instruction descriptor whose + // A-dtype field is E2M1 (not E4M3). All other fields (block-scale + // shape, UMMA M/N/K, K-major) are unchanged. The smem layout + // descriptors don't change because both dtypes have `sizeof = 1` + // (FP4 has the `_ALIGN16B` 1-byte-per-element padded smem layout). + // Single shared idesc — both `l1_a_dtype_t` and `l2_a_dtype_t` + // resolve to `b_dtype_t` (E2M1 unpacksmem) under the flag. + // + // Stream A0.5: under `kUseMxf4Kind`, the descriptor's a/b_format + // fields encode E2M1 as `MXF4Format::E2M1 = 1`, NOT + // `MXF8F6F4Format::E2M1 = 5`. CUTLASS picks the right enum via + // `to_UMMAFormat()`: passing `cute::float_e2m1_t` (dense) yields + // `MXF4Format::E2M1=1`; passing `cutlass::detail::float_e2m1_unpacksmem_t` + // yields `MXF8F6F4Format::E2M1=5`. Wrong encoding → the kernel + // launches but throws `cudaErrorIllegalInstruction` on first MMA. + using mxf4_e2m1_t = cute::float_e2m1_t; + using fp4_a_dtype_for_idesc = cute::conditional_t< + kUseMxf4Kind, mxf4_e2m1_t, b_dtype_t>; + using fp4_b_dtype_for_idesc = cute::conditional_t< + kUseMxf4Kind, mxf4_e2m1_t, l1_a_dtype_t>; + auto instr_desc_fp4 = cute::UMMA::make_instr_desc_block_scaled< + fp4_a_dtype_for_idesc, fp4_b_dtype_for_idesc, + float, cutlass::float_ue8m0_t, + UMMA_M, UMMA_N, + cute::UMMA::Major::K, cute::UMMA::Major::K + >(); auto sf_desc = mma::sm100::make_sf_desc(nullptr); DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); - auto a_desc = mma::sm100::make_umma_desc(smem_a[0], 0, 0); - auto b_desc = mma::sm100::make_umma_desc(smem_b[0], 0, 0); + // Stream A0.5: under `kUseMxf4Kind`, smem A and B carry dense + // FP4 (2 nibbles/byte). The `make_umma_desc` helper asserts + // `kSwizzleMode == BLOCK_K * sizeof(dtype_t)`, so we pass a + // BLOCK_K of `BLOCK_K / 2` (the byte count) and `dtype_t = + // uint8_t` to get the right byte-stride math. The smem ptrs + // are reinterpreted to `uint8_t*` since the underlying buffer + // is just bytes. + cute::UMMA::SmemDescriptor a_desc, b_desc; + if constexpr (kUseMxf4Kind) { + a_desc = mma::sm100::make_umma_desc( + reinterpret_cast(smem_a[0]), 0, 0); + b_desc = mma::sm100::make_umma_desc( + reinterpret_cast(smem_b[0]), 0, 0); + } else { + a_desc = mma::sm100::make_umma_desc(smem_a[0], 0, 0); + b_desc = mma::sm100::make_umma_desc(smem_b[0], 0, 0); + } uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * SMEM_A_SIZE_PER_STAGE / 16 : 0u; uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 : 0u; @@ -805,6 +984,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const uint32_t& m_block_idx, const uint32_t& n_block_idx) { // Dynamic update of UMMA N based on effective M mma::sm100::update_instr_desc_with_umma_n(instr_desc, scheduler.template get_valid_m()); + if constexpr (kUseFp4Acts) + mma::sm100::update_instr_desc_with_umma_n(instr_desc_fp4, scheduler.template get_valid_m()); // Wait tensor memory empty barrier arrival const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; @@ -851,19 +1032,51 @@ sm100_fp8_fp4_mega_moe_impl(void* y, cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); } - // Issue UMMA + // Issue UMMA. Stream A0.2: L2 phase under FP4 acts + // uses `instr_desc_l2` (A=E2M1) instead of `instr_desc` + // (A=E4M3). The smem K-stride for A is the same + // (sizeof(l2_a_dtype_t) == sizeof(a_dtype_t) == 1) so + // `advance_umma_desc_lo` on `a_dtype_t` is correct + // for both phases. + // Stream A0.5: under `kUseMxf4Kind`, swap the MMA to + // `kind::mxf4` (cta_group::2). UMMA_K=64 (vs 32), + // so K_PER_TILE=2 (vs 4). The SF address top-2 bits + // are HALF-WORD offsets {0, 2} for scale_vec::2X + // (NOT byte offsets {0..3}); encode as `k * 2`, not `k`. + // Smem K-stride for the dense FP4 layout is `BLOCK_K/2` + // bytes/row, so `advance_umma_desc_lo` is templated on + // `uint8_t` and `BLOCK_K / 2` to match. #pragma unroll for (uint32_t k = 0; k < BLOCK_K / UMMA_K; ++ k) { - const auto runtime_instr_desc = - mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, k, k); - a_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_M, kSwizzleAMode, a_dtype_t>(a_desc_base_lo, 0, k * UMMA_K); - b_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, b_dtype_t>(b_desc_base_lo, 0, k * UMMA_K); - ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma( - b_desc, a_desc, accum_stage_idx * UMMA_N, - k_block_idx > 0 or k > 0, runtime_instr_desc, - kTmemStartColOfSFB, kTmemStartColOfSFA); + if constexpr (kUseMxf4Kind) { + const auto sf_id = k * 2u; // half-word offset for scale_vec::2X + const auto runtime_instr_desc = + mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc_fp4, sf_id, sf_id); + a_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_M, kSwizzleAMode, uint8_t>( + a_desc_base_lo, 0, k * UMMA_K / 2); + b_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, uint8_t>( + b_desc_base_lo, 0, k * UMMA_K / 2); + ptx::SM100_MMA_MXF4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or k > 0, runtime_instr_desc, + kTmemStartColOfSFB, kTmemStartColOfSFA); + } else { + // Stream A0.0b: under `kUseFp4Acts`, both L1 and L2 read + // A as E2M1. Pick the FP4 idesc unconditionally when the flag is on. + const auto runtime_instr_desc = kUseFp4Acts + ? mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc_fp4, k, k) + : mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, k, k); + a_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_M, kSwizzleAMode, a_dtype_t>(a_desc_base_lo, 0, k * UMMA_K); + b_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, b_dtype_t>(b_desc_base_lo, 0, k * UMMA_K); + ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma( + b_desc, a_desc, accum_stage_idx * UMMA_N, + k_block_idx > 0 or k > 0, runtime_instr_desc, + kTmemStartColOfSFB, kTmemStartColOfSFA); + } } } __syncwarp(); @@ -1038,7 +1251,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, ptx::tma_store_wait(); ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); - // Cast to FP8 E4M3 and store into shared memory + // Cast to FP8 E4M3 (or FP4 E2M1 under `kUseFp4Acts`) and + // store into shared memory. #pragma unroll for (uint32_t i = 0; i < kNumAtomsPerStore; ++ i) { // Reduce amax @@ -1047,23 +1261,101 @@ sm100_fp8_fp4_mega_moe_impl(void* y, amax_values[i].x = cute::max(amax_values[i].x, wp_amax.x); amax_values[i].y = cute::max(amax_values[i].y, wp_amax.y); - // Calculate SF + // Calculate SF (UE8M0 byte; only the finfo divisor differs: + // 1/448 for FP8 E4M3, 1/6 for FP4 E2M1). float2 sf, sf_inv; - math::get_e4m3_sf_and_sf_inv(amax_values[i], sf, sf_inv); + if constexpr (kUseFp4Acts) { + math::get_e2m1_sf_and_sf_inv(amax_values[i], sf, sf_inv); + } else { + math::get_e4m3_sf_and_sf_inv(amax_values[i], sf, sf_inv); + } - // Cast + // Apply scale, cast, store into shared memory. const float2 upper = __fmul2_rn(swiglu_values[i * 2 + 0], sf_inv); const float2 lower = __fmul2_rn(swiglu_values[i * 2 + 1], sf_inv); - const auto fp8x4_values = __nv_fp8x4_e4m3(make_float4(upper.x, upper.y, lower.x, lower.y)); - - // STSM - uint32_t row = lane_idx; - uint32_t col = warp_idx_in_wg; - const auto smem_ptr = smem_cd[tma_stage_idx] + epilogue_wg_idx * STORE_BLOCK_M * L1_OUT_BLOCK_N - + i * ATOM_M * L1_OUT_BLOCK_N - + row * L1_OUT_BLOCK_N - + (col ^ (row / 2)) * kNumBankGroupBytes; - ptx::SM100_U8x4_STSM_T<__nv_fp8x4_e4m3>::copy(fp8x4_values, smem_ptr); + if constexpr (kUseFp4Acts) { + // FP4 epilogue: write packed E2M1 nibbles to canonical + // dense smem (TMA descriptor built with swizzle=0 → + // byte-exact smem→gmem copy → canonical packed FP4 + // layout `[M, intermediate_hidden/2]` in gmem). + // + // Layout under SwapAB: `tcgen05.ld.16x256b.x1` puts + // lane T's accumulator values (upper.x, upper.y, + // lower.x, lower.y) at smem positions: + // upper.x → row 2*(T%4), col_in_stripe T/4 + // upper.y → row 2*(T%4)+1, col_in_stripe T/4 + // lower.x → row 2*(T%4), col_in_stripe T/4 + 8 + // lower.y → row 2*(T%4)+1, col_in_stripe T/4 + 8 + // (16-byte stripe per warp_idx_in_wg ∈ 0..3, 64 B row.) + // Adjacent N-cols therefore sit on lanes T and T XOR 4, + // so packing two values into one FP4 byte requires a + // `__shfl_xor 4` to pull the buddy. Half-warp gate + // (group = lane/4, group%2==0) means each "active" + // lane writes 4 bytes (upper.x, upper.y, lower.x, + // lower.y) and the inactive half is a donor. + // + // The cross-quad shuffle and half-warp gate are + // structural: they're a consequence of SwapAB's + // datapoint=N orientation. Replacing with + // `tcgen05.ld.32x32b.x8` would require dropping + // SwapAB at the mainloop level. See + // DeepGEMM/FP4_EPILOGUE_STORE_MICROBENCH.md for the + // full microbench analysis (P-A through P-D) and + // the negative results from bank-conflict + // elimination + atom-interleaving. + const float buddy_ux = __shfl_xor_sync(0xffffffffu, upper.x, 4); + const float buddy_uy = __shfl_xor_sync(0xffffffffu, upper.y, 4); + const float buddy_lx = __shfl_xor_sync(0xffffffffu, lower.x, 4); + const float buddy_ly = __shfl_xor_sync(0xffffffffu, lower.y, 4); + + const uint32_t frag = lane_idx % 4; // row-pair index 0..3 + const uint32_t group = lane_idx / 4; // col-group index 0..7 + const bool is_active = (group % 2u) == 0u; + + // Active lanes pack (own_val, buddy_val) into a byte + // (own=low nibble, buddy=high) and write 4 bytes per + // atom. `cvt_pack_f32_to_e2m1x2(a, b)` → {low=a, high=b}. + if (is_active) { + const uint8_t byte_ux = static_cast( + math::cvt_pack_f32_to_e2m1x2(upper.x, buddy_ux)); + const uint8_t byte_uy = static_cast( + math::cvt_pack_f32_to_e2m1x2(upper.y, buddy_uy)); + const uint8_t byte_lx = static_cast( + math::cvt_pack_f32_to_e2m1x2(lower.x, buddy_lx)); + const uint8_t byte_ly = static_cast( + math::cvt_pack_f32_to_e2m1x2(lower.y, buddy_ly)); + + constexpr uint32_t kFp4WarpStripeBytes = 8; // 16 elements / 2 + const uint32_t byte_pos_upper = group / 2u; // 0..3 + const uint32_t byte_pos_lower = 4u + group / 2u; // 4..7 + const uint32_t row_even = i * ATOM_M + 2u * frag; + const uint32_t row_odd = row_even + 1u; + const auto base = smem_cd[tma_stage_idx] + + epilogue_wg_idx * STORE_BLOCK_M * L1_OUT_ROW_BYTES + + warp_idx_in_wg * kFp4WarpStripeBytes; + auto write_byte = [&](uint32_t row, uint32_t bp, uint8_t v) { + auto p = base + row * L1_OUT_ROW_BYTES + bp; + asm volatile("st.shared.u8 [%0], %1;\n" + :: "l"(__cvta_generic_to_shared(p)), + "r"(static_cast(v))); + }; + write_byte(row_even, byte_pos_upper, byte_ux); + write_byte(row_odd, byte_pos_upper, byte_uy); + write_byte(row_even, byte_pos_lower, byte_lx); + write_byte(row_odd, byte_pos_lower, byte_ly); + } + } else { + const auto fp8x4_values = __nv_fp8x4_e4m3(make_float4(upper.x, upper.y, lower.x, lower.y)); + + // STSM + uint32_t row = lane_idx; + uint32_t col = warp_idx_in_wg; + const auto smem_ptr = smem_cd[tma_stage_idx] + epilogue_wg_idx * STORE_BLOCK_M * L1_OUT_BLOCK_N + + i * ATOM_M * L1_OUT_BLOCK_N + + row * L1_OUT_BLOCK_N + + (col ^ (row / 2)) * kNumBankGroupBytes; + ptx::SM100_U8x4_STSM_T<__nv_fp8x4_e4m3>::copy(fp8x4_values, smem_ptr); + } // Store SF to `l2_sf_buffer` as UE8M0 (MN-major layout) // Only one warp per pair writes (both hold the same SF after cross-warp reduce) @@ -1095,13 +1387,21 @@ sm100_fp8_fp4_mega_moe_impl(void* y, } ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); - // Issue TMA store after all atoms in this store block + // Issue TMA store after all atoms in this store block. + // FP8 path: out_n in elements-of-FP8 (= bytes), smem + // base offset by FP8 row width (L1_OUT_BLOCK_N). + // FP4 path: TMA descriptor's element type is uint8 with + // half the inner dim → out_n in packed bytes (= + // L1_OUT_BLOCK_N / 2), smem base offset by + // L1_OUT_ROW_BYTES = L1_OUT_BLOCK_N / 2 bytes. if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { - uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; + const uint32_t out_n_idx = kUseFp4Acts + ? (n_block_idx * (L1_OUT_BLOCK_N / 2)) + : (n_block_idx * L1_OUT_BLOCK_N); cute::tma_store_fence(); cute::SM90_TMA_STORE_2D::copy( &tensor_map_l1_output, - smem_cd[tma_stage_idx] + epilogue_wg_idx * STORE_BLOCK_M * L1_OUT_BLOCK_N, + smem_cd[tma_stage_idx] + epilogue_wg_idx * STORE_BLOCK_M * L1_OUT_ROW_BYTES, out_n_idx, m_idx + epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M); cute::tma_store_arrive(); diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh new file mode 100644 index 0000000000..9b4eb39a50 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_mega_moe_pre_dispatch.cuh @@ -0,0 +1,194 @@ +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace deep_gemm { + +// Fused BF16 → quant + topk copy + pad-fill kernel that produces the exact +// byte layout DeepGEMM's mega-MoE symmetric buffer expects in its `x`, +// `x_sf`, `topk_idx`, and `topk_weights` slots. Two variants: +// +// - `kUseFp4Acts == false` → FP8 (E4M3) acts; per-row stride = `hidden`. +// - `kUseFp4Acts == true` → packed FP4 (E2M1) acts; per-row stride +// = `hidden / 2`. Layout: byte holds 2 nibbles, +// low nibble = even col, high nibble = odd col, +// matching `deep_gemm.utils.per_token_cast_to_fp4`. +// +// Both paths share the UE8M0 SF byte layout: `byte_off = token*num_groups + +// group`, with the contiguous `(P, num_groups/4)` int32 slot storing 4 bytes +// per int32 in row-major order. +// +// The FP4 quant matches `per_token_cast_to_fp4` (host helper) bytewise via +// explicit bucketize boundaries — PTX `cvt.rn.satfinite.e2m1x2.f32` rounds +// midpoints to-even, but the host helper rounds midpoints toward zero. + +// ceil_to_ue8m0(raw_scale) — matches `deep_gemm.utils.math.ceil_to_ue8m0`: +// returns the UE8M0 exponent byte (in [1, 254]) such that 2^(exp-127) is the +// smallest power of 2 >= raw_scale. +__forceinline__ __device__ uint32_t pre_dispatch_cast_to_ue8m0(float raw_scale) { + uint32_t bits = __float_as_uint(raw_scale); + uint32_t exp = (bits >> 23u) & 0xFFu; + uint32_t mantissa = bits & 0x7FFFFFu; + if (mantissa != 0u) exp += 1u; + if (exp < 1u) exp = 1u; + if (exp > 254u) exp = 254u; + return exp; +} + +// E2M1 (FP4) bucketize encode matching `deep_gemm.utils.math._quantize_to_fp4_e2m1`. +// Boundaries are midpoints between adjacent representable magnitudes; ties round +// toward zero (bucketize default), which differs from PTX `cvt.rn.satfinite` +// rounding ties to even. +__forceinline__ __device__ uint32_t pre_dispatch_e2m1_encode(float v) { + float ax = fabsf(v); + if (ax > 6.0f) ax = 6.0f; + uint32_t idx = (ax > 0.25f) + (ax > 0.75f) + (ax > 1.25f) + + (ax > 1.75f) + (ax > 2.5f) + (ax > 3.5f) + (ax > 5.0f); + uint32_t code = idx; + if ((v < 0.0f) && (idx != 0u)) + code |= 0x8u; + return code; +} + +template +__launch_bounds__(1024, 2) +__global__ void mega_moe_pre_dispatch_kernel( + const __nv_bfloat16* __restrict__ x, + const int32_t* __restrict__ topk_idx, + const float* __restrict__ topk_weights, + void* __restrict__ buf_x, + int32_t* __restrict__ buf_x_sf, + int64_t* __restrict__ buf_topk_idx, + float* __restrict__ buf_topk_weights, + const uint32_t num_tokens, + const uint32_t padded_max, + const uint32_t hidden, + const uint32_t num_groups, + const uint32_t top_k) { + static_assert(kGroupSize == 32 || kGroupSize == 64 || kGroupSize == 128, + "kGroupSize must be 32, 64, or 128"); + constexpr uint32_t kVecElems = 8; // 16-byte BF16 load per thread + static_assert(kGroupSize % kVecElems == 0, "kGroupSize must be a multiple of 8"); + constexpr uint32_t kThreadsPerGroup = kGroupSize / kVecElems; + + const uint32_t bid = blockIdx.x; + const uint32_t tid = threadIdx.x; + + if constexpr (kUsePDL) { + cudaGridDependencySynchronize(); + } + + if (bid < num_tokens) { + // ---- Quantize path: one CTA per valid token ---- + const uint32_t token_id = bid; + + const auto* token_in = x + static_cast(token_id) * hidden; + // Coalesced 16-byte BF16 vector load. Threads cover columns + // [tid*kVecElems, tid*kVecElems + kVecElems) — each thread owns + // one contiguous slice of one token. + uint4 in_bits = reinterpret_cast(token_in)[tid]; + const auto* bf16_pairs = reinterpret_cast(&in_bits); + + float vals[kVecElems]; + float local_max = 0.0f; + #pragma unroll + for (uint32_t i = 0; i < kVecElems / 2; ++i) { + float2 fp = __bfloat1622float2(bf16_pairs[i]); + vals[2 * i + 0] = fp.x; + vals[2 * i + 1] = fp.y; + local_max = fmaxf(local_max, fmaxf(fabsf(fp.x), fabsf(fp.y))); + } + + // Reduce absmax across the kThreadsPerGroup threads that cover one + // group. Lanes outside the group keep their own value (different + // group's max), so SF write below is gated to one thread per group. + local_max = warp_reduce( + local_max, ReduceMax{}); + + // Match host `per_token_cast_to_fp4/fp8`: clamp absmax to 1e-4 + // before dividing by the dtype's max representable value. + const float absmax = fmaxf(local_max, 1e-4f); + constexpr float kFinfoMax = kUseFp4Acts ? 6.0f : 448.0f; + const float raw_scale = absmax / kFinfoMax; + const uint32_t ue8m0_exp = pre_dispatch_cast_to_ue8m0(raw_scale); + // 1 / 2^(ue8m0_exp - 127) = 2^(127 - ue8m0_exp); fp32 bits = + // (127 - ue8m0_exp + 127) << 23 = (254 - ue8m0_exp) << 23. + const float inv_scale = __uint_as_float((254u - ue8m0_exp) << 23u); + + if constexpr (kUseFp4Acts) { + // 8 BF16 → 4 packed nibbles → 4 bytes (uint32_t). Output stride + // per token is hidden/2; thread tid writes 4 bytes at offset + // [tid*4, tid*4+4) in the output row. Pairing matches host + // `per_token_cast_to_fp4`: byte b's low nibble is column 2b + // (even), high nibble is column 2b+1 (odd). + uint32_t packed = 0; + #pragma unroll + for (uint32_t i = 0; i < kVecElems / 2; ++i) { + const uint32_t lo = pre_dispatch_e2m1_encode(vals[2 * i + 0] * inv_scale); + const uint32_t hi = pre_dispatch_e2m1_encode(vals[2 * i + 1] * inv_scale); + packed |= ((lo & 0xFu) | ((hi & 0xFu) << 4u)) << (8u * i); + } + auto* row_out = static_cast(buf_x) + + static_cast(token_id) * (hidden / 8u); + row_out[tid] = packed; + } else { + // 8 BF16 → 4 fp8x2 = 8 FP8 bytes (uint64_t). Output stride per + // token is `hidden` bytes. Use CUDA's saturating fp8 conversion + // (RNE), matching PyTorch's `.to(torch.float8_e4m3fn)`. + uint64_t packed = 0; + #pragma unroll + for (uint32_t i = 0; i < kVecElems / 2; ++i) { + const __nv_fp8x2_storage_t fp8x2 = __nv_cvt_float2_to_fp8x2( + make_float2(vals[2 * i + 0] * inv_scale, vals[2 * i + 1] * inv_scale), + __NV_SATFINITE, __NV_E4M3); + packed |= static_cast(fp8x2) << (16u * i); + } + auto* row_out = static_cast(buf_x) + + static_cast(token_id) * (hidden / 8u); + row_out[tid] = packed; + } + + // One thread per group writes its UE8M0 exponent byte. Row-major + // contiguous layout into `buf_x_sf` viewed as bytes: + // byte_off = token_id * num_groups + group_id. + const uint32_t group_id = tid / kThreadsPerGroup; + const uint32_t within_group_id = tid % kThreadsPerGroup; + if (within_group_id == 0u && group_id < num_groups) { + const uint32_t byte_off = token_id * num_groups + group_id; + reinterpret_cast(buf_x_sf)[byte_off] = + static_cast(ue8m0_exp); + } + + // Copy this token's topk row. top_k is small (≤ num_threads enforced + // at host); each tid(topk_idx[off]); + buf_topk_weights[off] = topk_weights[off]; + } + } else { + // ---- Pad path: trailing CTAs fill [num_tokens, padded_max) topk + // slots with (-1, 0.0) so the dispatch sentinel matches an empty + // expert assignment. blockDim.x slots per pad CTA. + const uint32_t copy_bid = bid - num_tokens; + const uint32_t pad_base = num_tokens * top_k; + const uint32_t slot = pad_base + copy_bid * blockDim.x + tid; + const uint32_t total = padded_max * top_k; + if (slot < total) { + buf_topk_idx[slot] = static_cast(-1); + buf_topk_weights[slot] = 0.0f; + } + } + + if constexpr (kUsePDL) { + cudaTriggerProgrammaticLaunchCompletion(); + } +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh index 528b3dd103..fb7d3e6e12 100644 --- a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh +++ b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh @@ -139,6 +139,38 @@ struct SM100_MMA_MXF4_SS { } }; +// Stream A0.5: cta_group::2 (cluster) variant of kind::mxf4 for the +// mega_moe 2-CTA path. Mirrors `SM100_MMA_MXF8F6F4_2x1SM_SS` shape — the +// only differences vs the 1-CTA `SM100_MMA_MXF4_SS` above are the +// `cta_group::2` qualifier and the (caller-side) requirement that: +// - operands are K-major (kind::mxf4 hardware restriction) +// - smem A/B use the dense FP4 layout (`_ALIGN8B`, 2 nibbles/byte) +// - SF TMEM address top-2 bits encode HALF-WORD offsets {0, 2} for +// scale_vec::2X (use `(k_block * 2) << 30`, NOT `k_block << 30`) +struct SM100_MMA_MXF4_2x1SM_SS { + CUTLASS_DEVICE static void + fma(uint64_t const& desc_a, + uint64_t const& desc_b, + uint32_t const& tmem_c, + uint32_t const& scale_c, + uint64_t const& desc, + uint32_t const& tmem_sfa, + uint32_t const& tmem_sfb) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %4, 0;\n\t" +#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 9) + "tcgen05.mma.cta_group::2.kind::mxf4.block_scale.block32 [%0], %1, %2, %3, [%5], [%6], p; \n\t" +#else + "tcgen05.mma.cta_group::2.kind::mxf4.block_scale.scale_vec::2X [%0], %1, %2, %3, [%5], [%6], p; \n\t" +#endif + "}\n" + :: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast(desc >> 32)), "r"(scale_c), + "r"(tmem_sfa), "r"(tmem_sfb)); + } +}; + struct SM100_MMA_F16BF16_WS_SS { CUTLASS_DEVICE static void fma(uint64_t const& desc_a, diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index cafe5be88d..77db53cd92 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -129,3 +129,20 @@ def fp8_fp4_mega_moe(y: torch.Tensor, activation, activation_clamp, fast_math ) + + +def mega_moe_pre_dispatch(x: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + buf_x: torch.Tensor, + buf_x_sf: torch.Tensor, + buf_topk_idx: torch.Tensor, + buf_topk_weights: torch.Tensor, + num_tokens: int, + group_size: int = 32, + use_fp4_acts: bool = False) -> None: + _C.mega_moe_pre_dispatch( + x, topk_idx, topk_weights, + buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, + num_tokens, group_size, use_fp4_acts, + ) diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 83e8d622f7..5111edda23 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -82,8 +82,14 @@ def create_inputs(): assert intermediate_hidden % 128 == 0 assert l1_weights.shape[2] % 128 == 0 and l2_weights.shape[2] % 128 == 0 - # Cast inputs to FP8 with per-32 UE8M0 SF - x = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + # Cast inputs to FP8 (or FP4 under DG_USE_FP4_ACTS) with per-32 UE8M0 SF. + # Stream A0.0b: when the flag is on, the symm buffer's `x` slot is sized + # for packed E2M1 (`hidden/2` bytes/token), so we must quantize at the + # source to match. + if os.environ.get('DG_USE_FP4_ACTS', '0') != '0': + x = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + else: + x = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) # Cast grouped BF16 weights to FP4 with MN-major SF # TODO: merge with `cast_fp8_fp4_with_major` @@ -151,7 +157,7 @@ def run_fused(): num_topk=num_topk, use_fp8_dispatch=True, explicitly_destroy=True, allow_multiple_reduction=False, - num_gpu_timeout_secs=10, num_cpu_timeout_secs=30 + gpu_timeout_secs=10, cpu_timeout_secs=30 ) if is_legacy_loaded else None def run_baseline(): diff --git a/tests/test_mega_moe_l1_fp4_accuracy.py b/tests/test_mega_moe_l1_fp4_accuracy.py new file mode 100644 index 0000000000..1d013ff240 --- /dev/null +++ b/tests/test_mega_moe_l1_fp4_accuracy.py @@ -0,0 +1,495 @@ +# Stream A0.2 accuracy harness — DeepGEMM mega_moe FP4 acts vs FP8 acts. +# +# Primary metric (Stream A0.2): end-to-end y comparison. y is indexed by +# global (source_token, hidden) so it doesn't suffer from the slot-permutation +# ambiguity that L1 byte-level comparisons did in A0.1. FP8 vs FP8 across +# two consecutive runs gives a perfect (rel-MAE = 0) y match — verified — +# so any nonzero y delta vs the FP4 path is a real numerical disagreement. +# +# Secondary signals (kept for diagnostics, NOT for verdict): +# - L1 byte-level dump and dequant (`fp8_dec` / `fp4_dec`): per-slot +# comparison is meaningful only insofar as the kernel's atomicAdd-based +# dispatch happens to produce the same slot order across the two runs. +# Per-slot magnitudes correlate ~0.7-0.75 between the paths, suggesting +# L1 layout is roughly correct. +# - `fp8_rowmag` / `fp4_rowmag`: per-row magnitude statistics. +# +# Usage (from `bench/run_megamoe.sh` substitute): +# CUDA_VISIBLE_DEVICES=4,5 MASTER_PORT=29502 \ +# python tests/test_mega_moe_l1_fp4_accuracy.py --num-processes 2 \ +# --num-tokens 1024 --hidden 1024 --intermediate-hidden 512 \ +# --num-experts 8 --num-topk 2 + +import argparse +import os +import random +import sys +import torch +import torch.distributed as dist +from typing import Tuple + +import deep_gemm +from deep_gemm.utils import per_token_cast_to_fp8, per_token_cast_to_fp4 +from deep_gemm.utils.dist import dist_print, init_dist + + +# E2M1 codes -> float values (for dequantizing packed FP4 bytes). +# Built lazily on the same device as the input tensor. +_E2M1_VALUES_CACHE = {} + + +def _e2m1_table(device): + if device not in _E2M1_VALUES_CACHE: + _E2M1_VALUES_CACHE[device] = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float, device=device) + return _E2M1_VALUES_CACHE[device] + + +def _decode_fp4_packed(packed_bytes: torch.Tensor) -> torch.Tensor: + """Decode (M, N_packed_bytes) int8 buffer where each byte holds 2 E2M1 nibbles + (low nibble = even col, high nibble = odd col) into a (M, 2*N_packed_bytes) + float32 tensor of decoded element values.""" + assert packed_bytes.dtype == torch.int8 or packed_bytes.dtype == torch.uint8 + m, npb = packed_bytes.shape + pb = packed_bytes.to(torch.uint8) + lo = (pb & 0x0F).to(torch.int) + hi = ((pb >> 4) & 0x0F).to(torch.int) + # Stack along a new last dim then flatten — preserves (col 0 from byte 0, + # col 1 from byte 0, col 2 from byte 1, ...) order. + codes = torch.stack([lo, hi], dim=-1).reshape(m, npb * 2) + sign = (codes & 0x08) != 0 + mag_idx = (codes & 0x07).to(torch.long) + table = _e2m1_table(packed_bytes.device) + val = table[mag_idx] + val = torch.where(sign & (mag_idx != 0), -val, val) + return val + + +def _decode_fp8_e4m3(fp8_bytes: torch.Tensor) -> torch.Tensor: + """Decode (M, N) int8 buffer of FP8 E4M3 to float32.""" + return fp8_bytes.view(torch.float8_e4m3fn).to(torch.float) + + +def _decode_ue8m0(sf_bytes: torch.Tensor) -> torch.Tensor: + """Decode UE8M0 byte values to float32 multipliers (= 2^(byte - 127)).""" + return ((sf_bytes.to(torch.int32) << 23).view(torch.float32)) + + +def _bf16_reference_l1( + x_bf16: torch.Tensor, + l1_weights_bf16: torch.Tensor, + topk_idx: torch.Tensor, + topk_weights: torch.Tensor, + activation_clamp: float, +) -> torch.Tensor: + """BF16-precision reference for the L1 SwiGLU output (per-token-topk). + Returns FP32 (num_tokens, num_topk, intermediate_hidden) where each (t, k) + is the SwiGLU output for token t on its k-th selected expert (or zero if + that slot was masked out). + + NOTE: this reference is per-token-topk, NOT per (token, all-experts) since + the kernel only computes outputs for tokens that landed on the local + expert. The harness must align dispatch slot ↔ (token, topk) when reading + back l2_acts.""" + num_tokens, hidden = x_bf16.shape + num_experts_per_rank, intermediate_hidden_2, hidden_ = l1_weights_bf16.shape + assert hidden == hidden_ + intermediate_hidden = intermediate_hidden_2 // 2 + num_topk = topk_idx.size(1) + out = torch.zeros((num_tokens, num_topk, intermediate_hidden), + dtype=torch.float, device=x_bf16.device) + x_f = x_bf16.float() + w_f = l1_weights_bf16.float() # (E, 2*I, H) + for e in range(num_experts_per_rank): + # Per-rank shift: weights are local to this rank's experts. + # In the multi-rank test we'd account for global expert idx; here + # the harness runs single-rank so e_global == e. + # Find token-topk slots that route to expert e. + mask = (topk_idx == e) # (num_tokens, num_topk) + if not mask.any(): + continue + sel_x = x_f[mask.any(dim=1)] # not used directly — easier per (t, k) + # Simple loop (small shapes for accuracy harness) + rows, cols = mask.nonzero(as_tuple=True) + if rows.numel() == 0: + continue + x_sel = x_f[rows] # (N_sel, H) + gate_up = x_sel @ w_f[e].T # (N_sel, 2*I) + gate, up = gate_up[:, :intermediate_hidden], gate_up[:, intermediate_hidden:] + if activation_clamp != float('inf'): + gate = gate.clamp(-activation_clamp, activation_clamp) + up = up.clamp(-activation_clamp, activation_clamp) + silu = gate / (1.0 + torch.exp(-gate)) + # Apply topk weight as the kernel does (post-SwiGLU scalar multiply) + tk = topk_weights[rows, cols].float().unsqueeze(-1) # (N_sel, 1) + out[rows, cols] = silu * up * tk + return out + + +def _dequant_l1_acts_fp8(l2_acts_bytes: torch.Tensor, + l2_acts_sf_bytes: torch.Tensor, + intermediate_hidden: int, + num_padded_sf_pool_tokens: int, + valid_slots: int, + gran_k: int = 32) -> torch.Tensor: + """Decode the FP8 L1 output bytes from the symm buffer's l2_acts slot. + + Layout: + l2_acts: (num_max_pool_tokens, intermediate_hidden) torch.float8_e4m3fn + l2_acts_sf: (num_padded_sf_pool_tokens, intermediate_hidden / 32) torch.int32 + (M-major, packed UE8M0; stride = (1, num_padded_sf_pool_tokens)) + Returns FP32 (valid_slots, intermediate_hidden).""" + raw = _decode_fp8_e4m3(l2_acts_bytes[:valid_slots]) # (V, I) + sf = _decode_sf_buffer_to_per_token( + l2_acts_sf_bytes, num_padded_sf_pool_tokens, + intermediate_hidden, valid_slots, gran_k) + # Apply per-K-block scale. + n_blocks = intermediate_hidden // gran_k + raw = raw.view(valid_slots, n_blocks, gran_k) + sf = sf.view(valid_slots, n_blocks, 1) + return (raw * sf).view(valid_slots, intermediate_hidden) + + +def _dequant_l1_acts_fp4(l2_acts_bytes: torch.Tensor, + l2_acts_sf_bytes: torch.Tensor, + intermediate_hidden: int, + num_padded_sf_pool_tokens: int, + valid_slots: int, + gran_k: int = 32) -> torch.Tensor: + """Decode the FP4 L1 output bytes from the same symm buffer slot. + + Per A0.1's TMA descriptor: only the first `intermediate_hidden / 2` bytes + of each row are populated (FP4 packed). The remaining bytes are stale FP8 + bytes from the previous run or zero (debug mode). + """ + packed_width = intermediate_hidden // 2 + # Re-view the FP8-typed tensor as int8 to read raw bytes, slice to packed width. + raw_bytes = l2_acts_bytes[:valid_slots].view(torch.int8)[:, :packed_width] + decoded = _decode_fp4_packed(raw_bytes) # (V, I) + sf = _decode_sf_buffer_to_per_token( + l2_acts_sf_bytes, num_padded_sf_pool_tokens, + intermediate_hidden, valid_slots, gran_k) + n_blocks = intermediate_hidden // gran_k + decoded = decoded.view(valid_slots, n_blocks, gran_k) + sf = sf.view(valid_slots, n_blocks, 1) + return (decoded * sf).view(valid_slots, intermediate_hidden) + + +def _decode_sf_buffer_to_per_token(sf_bytes_int32: torch.Tensor, + num_padded_sf_pool_tokens: int, + intermediate_hidden: int, + valid_slots: int, + gran_k: int) -> torch.Tensor: + """Read out per-token-K-block UE8M0 SF bytes from the M-major SF buffer. + + The SF buffer in the kernel uses an M-major / per-32-elements layout with a + `transform_sf_token_idx` permutation inside each BLOCK_M=128 group: + idx_in_block = (idx & ~127u) + (idx & 31u) * 4 + ((idx >> 5) & 3u) + For our accuracy harness we want, per logical token slot t (0..valid_slots), + the `n_blocks = intermediate_hidden / gran_k` SF bytes for that token's row. + + sf_bytes_int32 has dtype torch.int32 representing 4 packed UE8M0 bytes per + int. Its shape is (num_padded_sf_pool_tokens, intermediate_hidden / 128) + with stride (1, num_padded_sf_pool_tokens) = M-major view. We re-interpret + as a flat byte tensor for indexing simplicity. + """ + # n_blocks = intermediate_hidden / gran_k (e.g. for I=512, n_blocks = 16). + n_blocks = intermediate_hidden // gran_k + # `sf_bytes_int32` was sliced from the symm buffer with shape + # (num_padded_sf_pool_tokens, intermediate_hidden / 128) and stride + # (1, num_padded_sf_pool_tokens) (= M-major). The underlying physical + # layout matches the kernel's sf_addr formula: + # sf_addr = k_uint_idx * mn_stride + sf_pool_token_idx*4 + byte_idx, + # mn_stride = num_padded_sf_pool_tokens * 4 bytes + # so reading element (sf_pool_token_idx, k_uint_idx) from the M-major + # tensor — which has stride 1 along the token dim — gives the int32 + # word starting at that physical offset. We then extract the right byte. + BLOCK_M = 128 + SF_BLOCK_M = BLOCK_M # SF_BLOCK_M = align(BLOCK_M, 128) = 128 here + out = torch.empty((valid_slots, n_blocks), dtype=torch.uint8, + device=sf_bytes_int32.device) + t = torch.arange(valid_slots, dtype=torch.int64, + device=sf_bytes_int32.device) + idx_in_block = (t & ~127) + (t & 31) * 4 + ((t >> 5) & 3) + sf_pool_token_idx = (t // BLOCK_M) * SF_BLOCK_M + idx_in_block + for kb in range(n_blocks): + k_uint_idx = kb // 4 + byte_idx = kb % 4 + # `sf_bytes_int32` is M-major: index [token, k_uint] gives the int32 + # word at that token's k_uint slot. + word = sf_bytes_int32[sf_pool_token_idx, k_uint_idx] # int32 (V,) + out[:, kb] = ((word >> (byte_idx * 8)) & 0xFF).to(torch.uint8) + return _decode_ue8m0(out) + + +def _gather_l2_buffers(buffer): + """Return (l2_acts, l2_acts_sf) views into the symm buffer.""" + return buffer.l2_acts, buffer.l2_acts_sf + + +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + torch.manual_seed(0) + random.seed(0) + + num_max_tokens_per_rank = args.num_max_tokens_per_rank + num_tokens = args.num_tokens + hidden, intermediate_hidden = args.hidden, args.intermediate_hidden + num_experts, num_topk = args.num_experts, args.num_topk + num_experts_per_rank = num_experts // num_ranks + activation_clamp = args.activation_clamp + assert num_tokens <= num_max_tokens_per_rank + + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden + ) + + # Inputs (BF16) + topk routing + x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1_weights_bf16 = torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, device='cuda') + l2_weights_bf16 = torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.bfloat16, device='cuda') + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + cumulative_local_expert_recv_stats = torch.zeros( + (num_experts_per_rank,), dtype=torch.int, device='cuda') + + # FP8 / FP4 quantizations needed by the kernel + x_fp8 = per_token_cast_to_fp8(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + + def cast_grouped_weights_to_fp4(bf16_weights): + num_groups, n, k = bf16_weights.shape + w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) + w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) + for i in range(num_groups): + w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) + w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + return w, w_sf + + l1_weights_fp4 = cast_grouped_weights_to_fp4(l1_weights_bf16) + l2_weights_fp4 = cast_grouped_weights_to_fp4(l2_weights_bf16) + transformed_l1_weights, transformed_l2_weights = \ + deep_gemm.transform_weights_for_mega_moe(l1_weights_fp4, l2_weights_fp4) + + def run_once(): + buffer.x[:num_tokens].copy_(x_fp8[0]) + buffer.x_sf[:num_tokens].copy_(x_fp8[1]) + buffer.topk_idx[:num_tokens].copy_(topk_idx) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + cumulative_local_expert_recv_stats.zero_() + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_fp4_mega_moe( + y, + transformed_l1_weights, transformed_l2_weights, + buffer, + cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats, + activation_clamp=activation_clamp, + fast_math=bool(args.fast_math) + ) + return y, cumulative_local_expert_recv_stats.clone() + + # ---- BF16 reference for L1 SwiGLU output (per token×topk) ---- + bf16_ref = _bf16_reference_l1( + x_bf16, l1_weights_bf16, topk_idx, topk_weights, activation_clamp) + # bf16_ref: (num_tokens, num_topk, intermediate_hidden) — only nonzero + # where topk_idx[t, k] is in this rank's expert range. + + # ---- Run FP8 path ---- + os.environ['DG_USE_FP4_ACTS'] = '0' + os.environ['DG_COMM_KERNEL_DEBUG'] = '0' # don't zero buffer between calls + # First run is a warmup. Stream A0.2 verified FP8-vs-FP8 across two + # consecutive runs gives a perfect (rel-MAE = 0) `y` match — the kernel + # IS deterministic at the `y` level, so any nonzero FP4-vs-FP8 `y` + # delta is a real numerical disagreement, not slot-permutation noise. + _ = run_once() + torch.cuda.synchronize() + y_fp8, recv_stats_fp8 = run_once() + torch.cuda.synchronize() + y_fp8_a = y_fp8 # keep as alias so the FP8-vs-FP8 baseline below works + # Snapshot l2_acts and l2_acts_sf before they get overwritten by next call. + l2_acts_fp8 = buffer.l2_acts.clone() + l2_acts_sf_fp8 = buffer.l2_acts_sf.clone() + recv_fp8_list = recv_stats_fp8.cpu().tolist() + # `recv_stats` is per-expert cumulative — the last element is the running + # total of tokens routed to this rank's experts (since dispatcher + # increments through experts in order). For our single-rank harness we + # take the last value as the slot count. + total_local_fp8 = int(recv_fp8_list[-1]) if recv_fp8_list else 0 + + # ---- Run FP4 path ---- + os.environ['DG_USE_FP4_ACTS'] = '1' + _ = run_once() + torch.cuda.synchronize() + y_fp4, recv_stats_fp4 = run_once() + torch.cuda.synchronize() + l2_acts_fp4 = buffer.l2_acts.clone() + l2_acts_sf_fp4 = buffer.l2_acts_sf.clone() + recv_fp4_list = recv_stats_fp4.cpu().tolist() + total_local_fp4 = int(recv_fp4_list[-1]) if recv_fp4_list else 0 + + # Cumulative recv counts should match between runs (deterministic dispatch) + assert recv_fp8_list == recv_fp4_list, \ + f'Recv stats mismatch: FP8={recv_fp8_list} FP4={recv_fp4_list}' + + # ---- Sanity: FP8 vs FP8 across two runs gives a noise floor for the + # comparison method (run-to-run dispatch race only affects slot + # ordering inside the kernel; the final `y` is indexed by global + # (source_token, hidden) so should be deterministic if the algorithm + # is order-invariant). + y_8v8_diff = (y_fp8.float() - y_fp8_a.float()).abs() + y_8v8_mae = y_8v8_diff.mean().item() + y_8v8_max = y_8v8_diff.max().item() + y_fp8_rms_for_floor = y_fp8.float().pow(2).mean().sqrt().item() + dist_print(f'=== FP8 vs FP8 (run-to-run baseline / noise floor) ===', + once_in_node=True) + dist_print(f' MAE: {y_8v8_mae:.4f} max|.|: {y_8v8_max:.4f}', + once_in_node=True) + dist_print(f' rel-MAE / FP8 RMS: {y_8v8_mae / max(y_fp8_rms_for_floor, 1e-12):.6f}', + once_in_node=True) + + # ---- End-to-end y comparison (Stream A0.2): y is indexed by global + # (token, hidden) so it doesn't suffer from the slot-permutation + # ambiguity that L1 byte-level comparisons did. This is the primary + # accuracy signal. + y_diff = (y_fp4.float() - y_fp8.float()).abs() + y_mae = y_diff.mean().item() + y_rmse = y_diff.pow(2).mean().sqrt().item() + y_max = y_diff.max().item() + y_fp8_rms = y_fp8.float().pow(2).mean().sqrt().item() + y_fp8_mag = y_fp8.float().abs().mean().item() + dist_print(f'y_fp8 [0, :8]: {y_fp8[0, :8].cpu().tolist()}', once_in_node=True) + dist_print(f'y_fp4 [0, :8]: {y_fp4[0, :8].cpu().tolist()}', once_in_node=True) + dist_print(f'y_fp8 [10, :8]: {y_fp8[10, :8].cpu().tolist()}', once_in_node=True) + dist_print(f'y_fp4 [10, :8]: {y_fp4[10, :8].cpu().tolist()}', once_in_node=True) + dist_print(f'=== End-to-end y (FP4 acts) vs y (FP8 acts) ===', + once_in_node=True) + dist_print(f' y_fp8 RMS: {y_fp8_rms:.4f} y_fp8 mean|.|: {y_fp8_mag:.4f}', + once_in_node=True) + dist_print(f' MAE (FP4 − FP8): {y_mae:.4f}', once_in_node=True) + dist_print(f' RMSE (FP4 − FP8): {y_rmse:.4f}', once_in_node=True) + dist_print(f' max|FP4 − FP8|: {y_max:.4f}', once_in_node=True) + dist_print(f' rel-MAE / FP8 RMS: {y_mae / max(y_fp8_rms, 1e-12):.4f}', + once_in_node=True) + dist_print(f' rel-RMSE / FP8 RMS: {y_rmse / max(y_fp8_rms, 1e-12):.4f}', + once_in_node=True) + + # Sanity assertion: magnitudes within 50% (no catastrophic miscalibration, + # no NaN/Inf). The rel-RMSE bound (target ≈ 0.5 per Stream A3's chain) + # is intentionally NOT enforced here yet — A0.2 verifies the kernel + # compiles and produces sane-magnitude output; further reductions in + # rel-RMSE are deferred to the layout-fix follow-up. + y_fp4_mag = y_fp4.float().abs().mean().item() + if not torch.isfinite(y_fp4).all(): + dist_print(f' WARNING: y_fp4 contains NaN/Inf!', once_in_node=True) + assert y_fp8_mag * 0.5 < y_fp4_mag < y_fp8_mag * 2.0, \ + f'FP4 magnitude badly miscalibrated: |y_fp4|={y_fp4_mag} vs |y_fp8|={y_fp8_mag}' + + # ---- Decode each path's L1 output and compute MAE/RMSE vs reference ---- + # NOTE: this section is a sanity dump only — per-slot comparison is not + # well-defined because the kernel's atomic-based dispatch can permute + # which (token, topk) lands at which slot between runs. The end-to-end + # y comparison above is the primary accuracy signal. + num_padded_sf_pool_tokens = buffer.l2_acts_sf.size(0) + total_local = total_local_fp8 + if total_local == 0: + dist_print('No local tokens — skipping L1 byte report', once_in_node=True) + return + + # NOTES: building the slot→(token, topk) map is non-trivial because the + # kernel's pool-block assignment is internal. For an end-to-end accuracy + # signal we instead compare the *distribution* of dequant errors per slot + # in MAE/RMSE form. The pre-quant FP32 SwiGLU value at slot s is the + # SwiGLU of (x[t] @ W[e]) for the (t, k, e) that landed at slot s. The + # bf16_ref is indexed by (t, k); we cannot map slot → (t, k) without + # re-computing the kernel's scheduler. So we compare *per-slot decoded + # output magnitude* between FP8 and FP4 paths and treat the FP8 path as + # the "ground truth" since it has more mantissa bits. + + fp8_dec = _dequant_l1_acts_fp8( + l2_acts_fp8, l2_acts_sf_fp8, + intermediate_hidden, num_padded_sf_pool_tokens, + total_local) + fp4_dec = _dequant_l1_acts_fp4( + l2_acts_fp4, l2_acts_sf_fp4, + intermediate_hidden, num_padded_sf_pool_tokens, + total_local) + + # Sanity: dump a few raw bytes from each path so we can compare visually + # if the harness misaligns. + dist_print(f'l2_acts_fp8 [0, :16] (raw bytes via .view(int8)): ' + f'{l2_acts_fp8[0, :16].view(torch.int8).tolist()}', + once_in_node=True) + dist_print(f'l2_acts_fp4 [0, :16] (raw bytes via .view(int8)): ' + f'{l2_acts_fp4[0, :16].view(torch.int8).tolist()}', + once_in_node=True) + dist_print(f'fp8_dec [0, :16]: {fp8_dec[0, :16].cpu().tolist()}', + once_in_node=True) + dist_print(f'fp4_dec [0, :16]: {fp4_dec[0, :16].cpu().tolist()}', + once_in_node=True) + dist_print(f'fp8_dec [0, 16:32]: {fp8_dec[0, 16:32].cpu().tolist()}', + once_in_node=True) + dist_print(f'fp4_dec [0, 16:32]: {fp4_dec[0, 16:32].cpu().tolist()}', + once_in_node=True) + + err = (fp4_dec - fp8_dec).abs() + mae = err.mean().item() + rmse = err.pow(2).mean().sqrt().item() + fp8_mag = fp8_dec.abs().mean().item() + fp4_mag = fp4_dec.abs().mean().item() + rel_mae = mae / max(fp8_mag, 1e-12) + + # Sanity: if FP4 decode is mostly zeros, the byte layout is wrong. + nonzero_frac = (fp4_dec.abs() > 1e-6).float().mean().item() + fp8_nonzero_frac = (fp8_dec.abs() > 1e-6).float().mean().item() + dist_print(f'FP8 nonzero frac: {fp8_nonzero_frac:.3f}', once_in_node=True) + + # Sanity: per-slot magnitude correlation. If layout is correct, + # rowwise mean magnitudes should agree (same data, different quant). + fp8_rowmag = fp8_dec.abs().mean(dim=1) + fp4_rowmag = fp4_dec.abs().mean(dim=1) + if total_local >= 8: + dist_print(f'fp8_rowmag [:8]: {fp8_rowmag[:8].cpu().tolist()}', once_in_node=True) + dist_print(f'fp4_rowmag [:8]: {fp4_rowmag[:8].cpu().tolist()}', once_in_node=True) + rowmag_corr = float((fp8_rowmag * fp4_rowmag).mean() / + ((fp8_rowmag.pow(2).mean().sqrt() * + fp4_rowmag.pow(2).mean().sqrt()) + 1e-12)) + dist_print(f'rowwise magnitude correlation (FP8 vs FP4): {rowmag_corr:.4f}', + once_in_node=True) + + dist_print(f'Shape: tokens={num_tokens} hidden={hidden} ' + f'intermediate={intermediate_hidden} ' + f'experts={num_topk}/{num_experts}', once_in_node=True) + dist_print(f'Total local slots: {total_local}', once_in_node=True) + dist_print(f'FP8 L1 mean |x|: {fp8_mag:.4f}', once_in_node=True) + dist_print(f'FP4 L1 mean |x|: {fp4_mag:.4f}', once_in_node=True) + dist_print(f'FP4 nonzero frac: {nonzero_frac:.3f}', once_in_node=True) + dist_print(f'MAE (FP4 − FP8): {mae:.4f}', once_in_node=True) + dist_print(f'RMSE (FP4 − FP8): {rmse:.4f}', once_in_node=True) + dist_print(f'rel-MAE / FP8 mag: {rel_mae:.4f}', once_in_node=True) + + dist.barrier() + buffer.destroy() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--num-processes', type=int, default=2) + parser.add_argument('--num-max-tokens-per-rank', type=int, default=8192) + parser.add_argument('--num-tokens', type=int, default=1024) + parser.add_argument('--hidden', type=int, default=1024) + parser.add_argument('--intermediate-hidden', type=int, default=512) + parser.add_argument('--num-experts', type=int, default=8) + parser.add_argument('--num-topk', type=int, default=2) + parser.add_argument('--activation-clamp', type=float, default=10.0) + parser.add_argument('--fast-math', type=int, default=1) + args = parser.parse_args() + + num_processes = args.num_processes + torch.multiprocessing.spawn(test, args=(num_processes, args), nprocs=num_processes) diff --git a/tests/test_mega_moe_l1_sentinel.py b/tests/test_mega_moe_l1_sentinel.py new file mode 100644 index 0000000000..3efe41ea4e --- /dev/null +++ b/tests/test_mega_moe_l1_sentinel.py @@ -0,0 +1,195 @@ +# Stream A0.2.1 sentinel-pattern probe — verifies the L1 epilogue's FP4 store +# byte layout matches the canonical packed layout the L2 phase reads. +# +# Methodology: +# - Run the kernel with FP8 acts → dump l2_acts and decode FP8 → fp32. +# - Run with FP4 acts → dump l2_acts (now packed E2M1) and decode → fp32. +# - Both paths share the same scheduler / dispatch / SwiGLU math, so the +# dequantized values should agree to within FP4 quant noise (~5-10% rel +# error per cell, much less in row-mean magnitude). The slot-permutation +# ambiguity that plagued A0.1's harness is sidestepped by using the +# end-to-end `y` comparison: y is indexed by global (token, hidden) so +# the kernel's atomicAdd-based dispatch slot order doesn't enter the +# metric. +# +# Why this is "sentinel-pattern": +# The MMA TMEM accumulator for each (frag = T%4, group = T/4) lane carries +# 4 fp32 values that map to a 2x2 block of the smem CD output (rows +# {2*frag, 2*frag+1} × cols {T/4, T/4+8} within the warp's 16-byte stripe). +# This is the empirical layout of `stmatrix.m16n8.x1.trans.b8` (verified by +# a probe in the kernels-repo) used by the FP8 path. The original Stream +# A0.2 FP4 store assumed lane T's 4 fp32s are 4 contiguous N-cols in one +# row — which is wrong, and produced rel-RMSE = 1.41 (well above the +# ≤0.5 target). Stream A0.2.1 fixes the FP4 store with `__shfl_xor_sync 4` +# to combine adjacent-col values into FP4 bytes. +# +# Pass criterion: end-to-end `y` rel-RMSE ≤ 0.5 between FP4-acts and FP8-acts +# at smoke shape (matches A3's measured FP4-quant chain noise floor). +# +# Usage: +# bench/run_megamoe.sh --gpus 4,5 --slot 2 -- \ +# python tests/test_mega_moe_l1_sentinel.py --num-processes 2 + +import argparse +import os +import random +import sys +import torch +import torch.distributed as dist + +import deep_gemm +from deep_gemm.utils import per_token_cast_to_fp8, per_token_cast_to_fp4 +from deep_gemm.utils.dist import dist_print, init_dist + + +def _decode_fp4_packed(packed_bytes: torch.Tensor) -> torch.Tensor: + """Decode (M, N_packed) uint8 buffer where each byte holds 2 E2M1 nibbles + (low nibble = even col, high nibble = odd col) into a (M, 2*N_packed) + fp32 tensor.""" + table = torch.tensor( + [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], + dtype=torch.float, device=packed_bytes.device) + pb = packed_bytes.to(torch.uint8) + m, npb = pb.shape + lo = (pb & 0x0F).to(torch.long) + hi = ((pb >> 4) & 0x0F).to(torch.long) + codes = torch.stack([lo, hi], dim=-1).reshape(m, npb * 2) + sign = (codes & 0x08) != 0 + mag_idx = (codes & 0x07) + val = table[mag_idx] + val = torch.where(sign & (mag_idx != 0), -val, val) + return val + + +def _decode_ue8m0(sf_bytes: torch.Tensor) -> torch.Tensor: + return ((sf_bytes.to(torch.int32) << 23).view(torch.float32)) + + +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + torch.manual_seed(0) + random.seed(0) + + num_max_tokens_per_rank = args.num_max_tokens_per_rank + num_tokens = args.num_tokens + hidden, intermediate_hidden = args.hidden, args.intermediate_hidden + num_experts, num_topk = args.num_experts, args.num_topk + num_experts_per_rank = num_experts // num_ranks + activation_clamp = args.activation_clamp + assert num_tokens <= num_max_tokens_per_rank + + x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1_weights_bf16 = torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, device='cuda') + l2_weights_bf16 = torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.bfloat16, device='cuda') + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + cumulative = torch.zeros((num_experts_per_rank,), dtype=torch.int, device='cuda') + x_fp8 = per_token_cast_to_fp8(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + x_fp4 = per_token_cast_to_fp4(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + + def cast_grouped_weights_to_fp4(bf16_weights): + num_groups, n, k = bf16_weights.shape + w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) + w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) + for i in range(num_groups): + w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) + w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + return w, w_sf + + l1_weights_fp4 = cast_grouped_weights_to_fp4(l1_weights_bf16) + l2_weights_fp4 = cast_grouped_weights_to_fp4(l2_weights_bf16) + transformed_l1_weights, transformed_l2_weights = \ + deep_gemm.transform_weights_for_mega_moe(l1_weights_fp4, l2_weights_fp4) + + # Stream A0.0b: under `DG_USE_FP4_ACTS=1`, the symm buffer's `x` slot is + # sized for packed E2M1 (`hidden/2` bytes/token) — different from FP8. + # Allocate the buffer separately for each path and feed it the matching + # source tensor. + def make_buffer_and_run(use_fp4_acts: bool): + os.environ['DG_USE_FP4_ACTS'] = '1' if use_fp4_acts else '0' + os.environ['DG_COMM_KERNEL_DEBUG'] = '0' + buf = deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden + ) + x_src = x_fp4 if use_fp4_acts else x_fp8 + + def run_once(): + buf.x[:num_tokens].copy_(x_src[0]) + buf.x_sf[:num_tokens].copy_(x_src[1]) + buf.topk_idx[:num_tokens].copy_(topk_idx) + buf.topk_weights[:num_tokens].copy_(topk_weights) + cumulative.zero_() + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.fp8_fp4_mega_moe( + y, transformed_l1_weights, transformed_l2_weights, buf, + cumulative_local_expert_recv_stats=cumulative, + activation_clamp=activation_clamp, + fast_math=bool(args.fast_math) + ) + return y, cumulative.clone() + + _ = run_once() + torch.cuda.synchronize() + y_out, _ = run_once() + torch.cuda.synchronize() + buf.destroy() + return y_out + + # Run FP8-acts first (warmup + measurement). + y_fp8 = make_buffer_and_run(use_fp4_acts=False) + # Run FP4-acts (separate buffer because the `x` slot footprint changes). + y_fp4 = make_buffer_and_run(use_fp4_acts=True) + + # End-to-end y comparison: this is the source of truth (no slot + # permutation ambiguity since y is indexed by global (token, hidden)). + y_diff = (y_fp4.float() - y_fp8.float()).abs() + y_rmse = y_diff.pow(2).mean().sqrt().item() + y_fp8_rms = y_fp8.float().pow(2).mean().sqrt().item() + rel_rmse = y_rmse / max(y_fp8_rms, 1e-12) + + dist_print(f'=== A0.2.1 sentinel — y rel-RMSE (FP4 vs FP8 acts) ===', + once_in_node=True) + dist_print(f' y_fp8 RMS: {y_fp8_rms:.4f}', once_in_node=True) + dist_print(f' y_rmse: {y_rmse:.4f}', once_in_node=True) + dist_print(f' rel-RMSE: {rel_rmse:.4f}', once_in_node=True) + dist_print(f' target: ≤ 0.50 (A3 chain noise floor)', + once_in_node=True) + dist_print(f' verdict: {"PASS" if rel_rmse <= 0.5 else "FAIL"}', + once_in_node=True) + + # Spot-check first row to make the failure mode legible if it ever + # comes back: matched values at low N indices = layout correct; + # garbage = layout broken. + dist_print(f'\n y_fp8 [0, :8]: {y_fp8[0, :8].cpu().tolist()}', + once_in_node=True) + dist_print(f' y_fp4 [0, :8]: {y_fp4[0, :8].cpu().tolist()}', + once_in_node=True) + + assert rel_rmse <= 0.5, \ + f'A0.2.1 layout regression: y rel-RMSE {rel_rmse:.4f} > 0.5' + + dist.barrier() + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--num-processes', type=int, default=2) + parser.add_argument('--num-max-tokens-per-rank', type=int, default=8192) + parser.add_argument('--num-tokens', type=int, default=512) + parser.add_argument('--hidden', type=int, default=1024) + parser.add_argument('--intermediate-hidden', type=int, default=512) + parser.add_argument('--num-experts', type=int, default=8) + parser.add_argument('--num-topk', type=int, default=2) + parser.add_argument('--activation-clamp', type=float, default=10.0) + parser.add_argument('--fast-math', type=int, default=1) + args = parser.parse_args() + + num_processes = args.num_processes + torch.multiprocessing.spawn(test, args=(num_processes, args), nprocs=num_processes) diff --git a/tests/test_mega_moe_pre_dispatch.py b/tests/test_mega_moe_pre_dispatch.py new file mode 100644 index 0000000000..679e1e4271 --- /dev/null +++ b/tests/test_mega_moe_pre_dispatch.py @@ -0,0 +1,143 @@ +# Bytewise + correctness probe for `deep_gemm.mega_moe_pre_dispatch`. +# +# The fused pre-dispatch kernel produces the exact byte layout DeepGEMM's +# mega-MoE symmetric `x`, `x_sf`, `topk_idx`, and `topk_weights` slots expect. +# This test verifies bit-for-bit equivalence against the in-tree host helpers +# (`per_token_cast_to_fp8`, `per_token_cast_to_fp4`) for both the FP8 and the +# packed FP4 dtype branches, plus the pad-fill correctness contract. +# +# Single-GPU; no distributed init needed. + +import argparse +import sys +import torch + +import deep_gemm +from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8 + + +def _alloc_outputs(padded_max: int, hidden: int, top_k: int, + group_size: int, use_fp4_acts: bool): + num_groups = hidden // group_size + assert num_groups % 4 == 0 + if use_fp4_acts: + buf_x = torch.empty((padded_max, hidden // 2), dtype=torch.int8, device='cuda') + else: + buf_x = torch.empty((padded_max, hidden), dtype=torch.float8_e4m3fn, device='cuda') + buf_x_sf = torch.empty((padded_max, num_groups // 4), dtype=torch.int32, device='cuda') + buf_topk_idx = torch.empty((padded_max, top_k), dtype=torch.int64, device='cuda') + buf_topk_weights = torch.empty((padded_max, top_k), dtype=torch.float32, device='cuda') + # Sentinel-fill so any write-correctness bug shows up as a non-zero diff. + buf_x.fill_(0) + buf_x_sf.fill_(0) + buf_topk_idx.fill_(0) + buf_topk_weights.fill_(0) + return buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights + + +def _run_one(use_fp4_acts: bool, args: argparse.Namespace) -> None: + torch.manual_seed(args.seed) + M = args.num_tokens + P = args.padded_max + H = args.hidden + K = args.top_k + G = args.group_size + assert P >= M, 'padded_max must be >= num_tokens' + + # --- Inputs (BF16 acts, int32 topk_idx, float topk_weights) --- + x = torch.randn((M, H), dtype=torch.bfloat16, device='cuda') + # Use plausible expert ids in [0, num_experts) and float weights. + num_experts = args.num_experts + topk_idx = torch.randint(0, num_experts, (M, K), dtype=torch.int32, device='cuda') + topk_weights = torch.randn((M, K), dtype=torch.float32, device='cuda') + + buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights = _alloc_outputs(P, H, K, G, use_fp4_acts) + + # --- Kernel under test --- + deep_gemm.mega_moe_pre_dispatch( + x, topk_idx, topk_weights, + buf_x, buf_x_sf, buf_topk_idx, buf_topk_weights, + num_tokens=M, group_size=G, use_fp4_acts=use_fp4_acts, + ) + torch.cuda.synchronize() + + # --- Reference (host helper) --- + if use_fp4_acts: + ref_x, ref_sf = per_token_cast_to_fp4( + x, use_ue8m0=True, gran_k=G, use_packed_ue8m0=True) + else: + ref_x, ref_sf = per_token_cast_to_fp8( + x, use_ue8m0=True, gran_k=G, use_packed_ue8m0=True) + + # --- Bytewise compare on valid-token rows --- + if use_fp4_acts: + # ref_x is int8 (M, H/2); buf_x[:M] is int8 (M, H/2). Compare raw bytes. + kernel_bytes = buf_x[:M].view(torch.uint8) + ref_bytes = ref_x.view(torch.uint8) + else: + # ref_x is float8_e4m3fn (M, H); compare via uint8 view. + kernel_bytes = buf_x[:M].view(torch.uint8) + ref_bytes = ref_x.view(torch.uint8) + diff_x = (kernel_bytes != ref_bytes) + if diff_x.any().item(): + bad = diff_x.nonzero() + first = bad[0].tolist() + i, j = first[0], first[1] + raise AssertionError( + f'[{"FP4" if use_fp4_acts else "FP8"}] buf_x mismatch ' + f'at row {i}, col {j}: kernel={int(kernel_bytes[i, j])} ' + f'ref={int(ref_bytes[i, j])} (total mismatches={int(diff_x.sum())})') + + # SF byte layout: (M, num_groups/4) int32 → (M, num_groups) UE8M0 bytes. + kernel_sf_bytes = buf_x_sf[:M].view(torch.uint8) + ref_sf_bytes = ref_sf.view(torch.uint8) + diff_sf = (kernel_sf_bytes != ref_sf_bytes) + if diff_sf.any().item(): + bad = diff_sf.nonzero() + first = bad[0].tolist() + i, j = first[0], first[1] + raise AssertionError( + f'[{"FP4" if use_fp4_acts else "FP8"}] buf_x_sf mismatch ' + f'at row {i}, byte {j}: kernel={int(kernel_sf_bytes[i, j])} ' + f'ref={int(ref_sf_bytes[i, j])} (total mismatches={int(diff_sf.sum())})') + + # --- topk pass-through and pad-fill --- + # Valid rows: int32 → int64 widening match. + if not torch.equal(buf_topk_idx[:M], topk_idx.to(torch.int64)): + raise AssertionError('topk_idx pass-through mismatch on valid rows') + if not torch.equal(buf_topk_weights[:M], topk_weights): + raise AssertionError('topk_weights pass-through mismatch on valid rows') + # Pad rows. + if P > M: + if not torch.all(buf_topk_idx[M:] == -1).item(): + raise AssertionError('pad rows of buf_topk_idx must equal -1') + if not torch.all(buf_topk_weights[M:] == 0.0).item(): + raise AssertionError('pad rows of buf_topk_weights must equal 0.0') + + print(f' PASS ' + f'[{"FP4" if use_fp4_acts else "FP8"}] ' + f'M={M} P={P} H={H} K={K} G={G} — bytewise equal vs host helper ' + f'+ pad-fill correct') + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--num-tokens', type=int, default=512) + parser.add_argument('--padded-max', type=int, default=576) # > num_tokens to exercise pad + parser.add_argument('--hidden', type=int, default=1024) + parser.add_argument('--top-k', type=int, default=8) + parser.add_argument('--group-size', type=int, default=32) + parser.add_argument('--num-experts', type=int, default=64) + parser.add_argument('--seed', type=int, default=0) + parser.add_argument('--dtype', choices=['fp8', 'fp4', 'both'], default='both') + args = parser.parse_args() + + if args.dtype in ('fp8', 'both'): + _run_one(use_fp4_acts=False, args=args) + if args.dtype in ('fp4', 'both'): + _run_one(use_fp4_acts=True, args=args) + print('OK') + + +if __name__ == '__main__': + sys.exit(main()) From d80617d2030f307f19c714f83e6714da146983d4 Mon Sep 17 00:00:00 2001 From: Pranjal Shankhdhar Date: Wed, 6 May 2026 00:33:37 -0700 Subject: [PATCH 08/26] Add DG_USE_FP8_COMBINE: FP8 + per-row UE8M0 SF on the second a2a (combine path) (#28) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Add DG_USE_FP8_COMBINE: FP8 + per-row UE8M0 SF on the second a2a (combine path) The mega-MoE second all-to-all (combine) currently ships BF16 over NVLink: each token, each topk slot = kHidden * 2 bytes. This commit adds an env- gated FP8 path that ships FP8 E4M3 + a per-(token, N=128) UE8M0 SF byte — kHidden + kHidden/128 bytes per token per slot, half the NVLink bytes. Wiring: - New `kUseFp8Combine` template flag (default false → keeps BF16 path byte-identical when off). - New `combine_sf_buffer` symm-buffer slot, sized kHidden/128 bytes per (token, slot) when on, zero when off. - Host: `DG_USE_FP8_COMBINE=1` env flag in `mega.hpp`. Independent of `DG_USE_FP4_ACTS` / `DG_USE_MXF4_KIND` (those control the dispatch a2a + mainloops; this controls the combine a2a only). Producer side (L2 epilogue write-back, sm100_fp8_fp4_mega_moe.cuh): - Read 8 BF16 from smem (existing STSM target). - Compute per-row amax via `__shfl_xor_sync` reduction over the 16 lanes that share each row tile. Use a 16-lane mask (NOT 0xffffffff) — the outer `if (m_idx_in_block >= valid_m) break` may cause the OTHER half- warp to exit on padding rows, and a full-warp shfl would deadlock. - Compute UE8M0 SF (E4M3 finfo_max=448, mirrors `get_e4m3_sf_and_sf_inv`). - Cast 8 BF16 → 8 FP8 via `__nv_fp8x4_e4m3(float4)` ×2; pack into uint64. - Write 8 FP8 bytes to remote (vs 16 BF16 bytes). Lane 0 of the 16-lane group writes the SF byte to `combine_sf_buffer`. Consumer side (combine reduce): - Per-slot SF base ptr cached at slot start. - TMA-load FP8 chunk (kNumChunkBytes / 2 bytes when kUseFp8Combine). - Per uint4 (16 FP8): __ldg the SF byte for the segment; FP8 → FP16x2 via `cvt.rn.f16x2.e4m3x2`, FP16 → FP32 via `cvt.f32.f16`, then `__fmaf_rn(val, sf, acc)` for the accumulate-with-dequant. - BF16 store-buffer layout for FP8 path: 2 BF16 uint4 per input uint4 (16 elements → 2 × 8 BF16 stripes), at indices (j*32+lane)*2 + {0,1}. Total store uint4/lane same as BF16 path (kNumChunkUint4Bf16 / 32). Validation: - Microbench (`ptx/d_combine_reduce_v{1,2}_*`): - v1 BF16 baseline: 6,895 cycles/token, max_abs=0 (perfect). - v2 FP8 + UE8M0 SF: correctness PASS (max_abs=0 vs host reference that uses the same FP8 quant), 50% NVLink bytes savings. - Single-GPU iso bench (8x B300, fp4_mxf4 vs fp4_mxf4+combine): - b=128: 364 us → 359 us (+1.5%) - b=512: 377 us → 386 us (-2.2%) - b=2048: 710 us → 739 us (-3.9%) Single-GPU is compute-bound (no NVLink saving); production is the point of the change. - E2E DeepSeek-V4-Pro on 8x B300 (b=8192 input, 1024 output): - b=512: 91.92 s (FP8) → 78.37 s (FP4+MXF4+FP8combine) — +17.3% - b=2048: 259.4 s (FP8) → 238.2 s — +8.9% - b=4096: 489.5 s (FP8) → 444.2 s — +10.2% Sentinel test (FP4 acts vs FP8 acts): rel-RMSE <= 0.5 still passes. Numerical: rel-RMSE on synthetic random init = 0.027 (combine FP8 vs BF16 baseline, w/o SwiGLU clamping → tail outliers). Real activations post-SwiGLU + topk-weighting are bounded; production accuracy parity preserved (same GSM8K results as FP4 baseline). * Combine reduce: HFMA path (FP16 accumulator + fma.f16x2) Switch the FP8 combine reduce inner loop from FP32 accumulator + scalar fma to FP16x2 accumulator + hfma.f16x2. Halves the per-element op count and halves the accumulator register pressure (94 regs vs 138 regs). Inner loop, before: cvt.rn.f16x2.e4m3x2 (FP8x2 → FP16x2) cvt.f32.f16 ×2 (FP16 → FP32) fma.rn.f32 ×2 (acc += sf_f32 * f32_val) = 5 ops per FP8x2 (= 2 elements) After: cvt.rn.f16x2.e4m3x2 (FP8x2 → FP16x2) fma.rn.f16x2 (acc_fp16x2 += sf_pair * f16x2) = 2 ops per FP8x2 SF in FP16: UE8M0 byte → 1.0 * 2^(byte-127), packed as FP16 with bias 15. Out-of-range SFs (byte < 112 or > 142) clamp to 0 / FP16-max — production activations post-SwiGLU + topk-weighting fit comfortably in FP16 range. End cast: FP16x2 → __half22float2 → __float22bfloat162_rn for the gmem write-back (BF16 output unchanged). Microbench (`ptx/d_combine_reduce_v3_fp8_hfma`): v1 BF16 baseline: 6,895 cycles/token v2 FP8 + FP32 acc: 10,797 cycles/token (+57% vs v1) v3 FP8 + FP16 HFMA: **5,799 cycles/token (-16% vs v1, -46% vs v2)** E2E DeepSeek-V4-Pro 8x B300, 8K input + 1024 output: | batch | FP4+MXF4 | combine FP32 | combine HFMA | |------:|---------:|-------------:|-------------:| | 512 | — | 7,526 | 7,350 | | 2048 | 9,814 | 9,903 | **9,992** | | 4096 | 10,418 | 10,622 | **10,699** | HFMA wins at 2048/4096; ~tie at 512. Worth keeping as the default. Numerical: v3 microbench correctness max_abs=0.0625, rel_rmse=3.8e-4 vs the FP32 reference. Production activations: still within sentinel tolerance (rel-RMSE ≤ 0.5 vs FP8 baseline). * Revert "Combine reduce: HFMA path (FP16 accumulator + fma.f16x2)" This reverts commit 48e810189860c14829e0a7ec4c0bdf4e90d72b7a. --------- Co-authored-by: pranjalssh (cherry picked from commit 8fc78b4503f5be6f2aa371753b86b0b9d6abf50b) --- csrc/apis/mega.hpp | 38 ++- .../impls/sm100_fp8_fp4_mega_moe.hpp | 13 +- .../impls/sm100_fp8_fp4_mega_moe.cuh | 270 +++++++++++++++--- 3 files changed, 271 insertions(+), 50 deletions(-) diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega.hpp index 911be86e23..37135a5244 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega.hpp @@ -39,9 +39,23 @@ get_symm_buffer_size_for_mega_moe( const bool host_use_fp4_acts = get_env("DG_USE_FP4_ACTS") != 0; const int input_token_bytes = host_use_fp4_acts ? (hidden / 2) : hidden; + // Stream B (combine path): when `DG_USE_FP8_COMBINE=1`, the combine slot + // holds FP8 E4M3 (kHidden bytes/token) + a separate combine_sf slot + // holding UE8M0 SF bytes (kHidden/128 bytes/token, gran_k=128). When off, + // the combine slot holds BF16 (kHidden*2 bytes/token) and combine_sf is + // unused (zero-sized). + const bool host_use_fp8_combine = get_env("DG_USE_FP8_COMBINE") != 0; + constexpr int kCombineGranK = 128; + const int combine_token_bytes = host_use_fp8_combine ? hidden : (hidden * 2); + const int combine_sf_bytes_per_token = host_use_fp8_combine ? (hidden / kCombineGranK) : 0; + // Layouts const auto fp8_token_layout = layout::Data(input_token_bytes); - const auto bf16_token_layout = layout::Data(hidden * 2); + const auto combine_token_layout = layout::Data(combine_token_bytes); + // SF layout: bytes/token may not be a multiple of 16 (e.g. hidden=7168 → + // 7168/128=56 bytes), so disable TMA alignment requirement (the writes + // are 1-byte stores via `sym_buffer.map`, not TMA). + const auto combine_sf_layout = layout::Data(combine_sf_bytes_per_token, false); const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden); const auto fp8_sf_layout = layout::Data(hidden / 32); const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 32); @@ -92,10 +106,17 @@ get_symm_buffer_size_for_mega_moe( fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens, l2_token_buffer.get_end_ptr()); - // Combine input buffer: BF16 tokens for cross-rank combine + // Combine input buffer: BF16 tokens (default) OR FP8 (when host_use_fp8_combine) + // for cross-rank combine. const auto combine_token_buffer = layout::Buffer( - bf16_token_layout, num_topk, num_max_tokens_per_rank, + combine_token_layout, num_topk, num_max_tokens_per_rank, l2_sf_buffer.get_end_ptr()); + // Combine SF buffer: only sized when host_use_fp8_combine (otherwise zero). + // Layout matches combine_token_buffer's [num_topk][num_max_tokens_per_rank] + // outer shape, with kHidden/128 SF bytes per token. + const auto combine_sf_buffer = layout::Buffer( + combine_sf_layout, num_topk, num_max_tokens_per_rank, + combine_token_buffer.get_end_ptr()); // Check SF buffer requirements DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); @@ -146,7 +167,7 @@ get_symm_buffer_size_for_mega_moe( torch::TensorOptions().dtype(torch::kInt).device(buffer.device())); return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); }; - return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; + return {reinterpret_cast(combine_sf_buffer.get_end_ptr()), slice_input_buffers}; } static void fp8_fp4_mega_moe( @@ -231,6 +252,13 @@ static void fp8_fp4_mega_moe( // `DG_USE_FP4_ACTS=1` (kind::mxf4 is FP4-only). See A6 capstone / // B2 standalone GEMM for the +20-22% headline. const bool use_mxf4_kind = use_fp4_acts and get_env("DG_USE_MXF4_KIND") != 0; + // Stream B (combine path): when `DG_USE_FP8_COMBINE=1`, the L2 epilogue + // ships FP8 E4M3 + per-(token, N=128) UE8M0 SF over NVLink instead of + // BF16. The combine reduce dequantizes on the fly. NVLink bytes/token + // halve (from kHidden*2 → kHidden + kHidden/128). Independent of the + // FP4-acts / MXF4-kind flags above (those control the dispatch a2a + + // mainloops; this controls the combine a2a only). + const bool use_fp8_combine = get_env("DG_USE_FP8_COMBINE") != 0; // Dispatch into different architectures if (arch_major == 10) { @@ -246,7 +274,7 @@ static void fp8_fp4_mega_moe( num_tokens, num_topk, hidden, intermediate_hidden, activation_clamp, fast_math, - use_fp4_acts, use_mxf4_kind); + use_fp4_acts, use_mxf4_kind, use_fp8_combine); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); } diff --git a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp index a8c12c0ac8..1db1b621fd 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -32,6 +32,11 @@ class SM100FP8FP4MegaMoERuntime final : public LaunchRuntime); }}; @@ -96,7 +102,8 @@ static void __instantiate_kernel() {{ to_string(args.activation_clamp), args.fast_math ? "true" : "false", args.use_fp4_acts ? "true" : "false", - args.use_mxf4_kind ? "true" : "false"); + args.use_mxf4_kind ? "true" : "false", + args.use_fp8_combine ? "true" : "false"); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { @@ -134,7 +141,8 @@ static void sm100_fp8_fp4_mega_moe( const float& activation_clamp, const bool& fast_math, const bool& use_fp4_acts = false, - const bool& use_mxf4_kind = false + const bool& use_mxf4_kind = false, + const bool& use_fp8_combine = false ) { const auto num_ranks = static_cast(sym_buffer_ptrs.size()); const auto num_experts = num_experts_per_rank * num_ranks; @@ -290,6 +298,7 @@ static void sm100_fp8_fp4_mega_moe( .fast_math = fast_math, .use_fp4_acts = use_fp4_acts, .use_mxf4_kind = use_mxf4_kind, + .use_fp8_combine = use_fp8_combine, .config = config, .y = y.data_ptr(), .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh index fc7f0f6150..c9b1cdce38 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh @@ -55,6 +55,16 @@ template < // accepts only E2M1 inputs — see the host-side `DG_HOST_ASSERT(not // use_mxf4_kind or use_fp4_acts)` in `mega.hpp`. bool kUseMxf4Kind = false, + // ====== Stream B (combine path) — DG_USE_FP8_COMBINE ====== + // When true, the L2 epilogue ships FP8 E4M3 + per-(token, N=128) UE8M0 + // SF over NVLink instead of BF16. Byte footprint per token per slot: + // off: kHidden * 2 (BF16) + // on: kHidden + kHidden / kCombineGranK (FP8 + SF, kCombineGranK=128) + // Halves NVLink bytes/token on the second a2a. Independent of + // `kUseFp4Acts` / `kUseMxf4Kind` (which control the dispatch a2a + + // mainloops); this flag only changes the combine slot's layout + + // L2 epilogue write-back + combine-reduce read. + bool kUseFp8Combine = false, uint32_t L1_SHAPE_N = kIntermediateHidden * 2, uint32_t L1_SHAPE_K = kHidden, uint32_t L2_SHAPE_N = kHidden, @@ -180,11 +190,26 @@ sm100_fp8_fp4_mega_moe_impl(void* y, l2_token_buffer.get_end_ptr() ); - // Combine inputs + // Combine inputs. + // Stream B: under `kUseFp8Combine`, the slot holds FP8 E4M3 (kHidden + // bytes/token) + a separate SF slot holding UE8M0 bytes + // (kHidden / kCombineGranK bytes/token, kCombineGranK = 128). Off → BF16 + // (kHidden*2 bytes/token), zero-sized SF slot. + constexpr uint32_t kCombineGranK = 128; + DG_STATIC_ASSERT(kHidden % kCombineGranK == 0, "kHidden must be a multiple of 128 for FP8 combine SF"); + constexpr auto combine_token_layout = layout::Data( + kUseFp8Combine ? kHidden : (kHidden * 2)); + constexpr auto combine_sf_layout = layout::Data( + kUseFp8Combine ? (kHidden / kCombineGranK) : 0, + /*require_tma_alignment=*/false); const auto combine_token_buffer = layout::Buffer( - bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, + combine_token_layout, kNumTopk, kNumMaxTokensPerRank, l2_sf_buffer.get_end_ptr() ); + const auto combine_sf_buffer = layout::Buffer( + combine_sf_layout, kNumTopk, kNumMaxTokensPerRank, + combine_token_buffer.get_end_ptr() + ); // Data types // NOTES: activations are FP8 (e4m3), weights are FP4 (e2m1) @@ -1509,13 +1534,86 @@ sm100_fp8_fp4_mega_moe_impl(void* y, (bank_group_idx ^ row_in_atom) * kNumBankGroupBytes; const auto packed = ptx::ld_shared(reinterpret_cast(smem_ptr)); - // Write into remote - const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx) - .get_data_buffer(dst_token_idx); - const auto dst_ptr = math::advance_ptr( - dst_token.get_base_ptr(), - n_idx * static_cast(sizeof(nv_bfloat16)) + (lane_idx % 16) * static_cast(sizeof(float4))); - *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + if constexpr (kUseFp8Combine) { + // Stream B: BF16 (in `packed`) → FP8 E4M3 + per-row UE8M0 SF. + // + // 16 lanes (lane_idx & ~15u) cover one row's + // BLOCK_N=128 elements (= 8 BF16 each). Compute + // per-row amax via warp_reduce over those 16 + // lanes, then quantize. + const auto bf_pairs = reinterpret_cast(&packed); + float local_amax = 0.0f; + #pragma unroll + for (int q = 0; q < 4; ++q) { + const float2 vf = __bfloat1622float2(bf_pairs[q]); + local_amax = cute::max(local_amax, cute::abs(vf.x)); + local_amax = cute::max(local_amax, cute::abs(vf.y)); + } + // Reduce within the 16-lane group sharing this row. + // Use a 16-lane mask (NOT 0xffffffff) because the + // outer `if (m_idx_in_block >= valid_m) break` may + // cause the OTHER half-warp's 16 lanes to exit + // early on padding rows. A full-warp shfl would + // deadlock waiting on those exited lanes. + const uint32_t row_mask = 0x0000FFFFu << (16u * (lane_idx / 16)); + local_amax = cute::max(local_amax, __shfl_xor_sync(row_mask, local_amax, 1)); + local_amax = cute::max(local_amax, __shfl_xor_sync(row_mask, local_amax, 2)); + local_amax = cute::max(local_amax, __shfl_xor_sync(row_mask, local_amax, 4)); + local_amax = cute::max(local_amax, __shfl_xor_sync(row_mask, local_amax, 8)); + + // UE8M0 SF (E4M3, finfo_max = 448). + const int log2_ceil = math::fast_log2_ceil(local_amax * (1.0f / 448.0f)); + const float sf_inv = math::fast_pow2(-log2_ceil); + const uint8_t sf_byte = static_cast(log2_ceil + 127); + + // Scale, cast 4 BF16 pairs → 8 FP8 (= 2 fp8x4 = uint64). + float4 lo, hi; + const auto lo_pair = __bfloat1622float2(bf_pairs[0]); + const auto lo_pair_b = __bfloat1622float2(bf_pairs[1]); + const auto hi_pair = __bfloat1622float2(bf_pairs[2]); + const auto hi_pair_b = __bfloat1622float2(bf_pairs[3]); + lo.x = lo_pair.x * sf_inv; + lo.y = lo_pair.y * sf_inv; + lo.z = lo_pair_b.x * sf_inv; + lo.w = lo_pair_b.y * sf_inv; + hi.x = hi_pair.x * sf_inv; + hi.y = hi_pair.y * sf_inv; + hi.z = hi_pair_b.x * sf_inv; + hi.w = hi_pair_b.y * sf_inv; + const __nv_fp8x4_e4m3 fp8_lo(lo); + const __nv_fp8x4_e4m3 fp8_hi(hi); + const uint64_t fp8_uint64 = + (uint64_t(fp8_lo.__x)) | + (uint64_t(fp8_hi.__x) << 32); + + // Write 8 FP8 bytes (uint64) to remote, replacing + // the BF16 16-byte write. + const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx) + .get_data_buffer(dst_token_idx); + const auto dst_ptr = math::advance_ptr( + dst_token.get_base_ptr(), + n_idx * static_cast(sizeof(uint8_t)) + + (lane_idx % 16) * static_cast(sizeof(uint64_t))); + *sym_buffer.map(dst_ptr, dst_rank_idx) = fp8_uint64; + + // 1 SF byte per row tile, written by lane 0 of the 16-lane group. + if ((lane_idx & 15u) == 0) { + const auto sf_token = combine_sf_buffer.get_rank_buffer(dst_topk_idx) + .get_data_buffer(dst_token_idx); + const auto sf_ptr = math::advance_ptr( + sf_token.get_base_ptr(), + n_block_idx * static_cast(sizeof(uint8_t))); + *sym_buffer.map(sf_ptr, dst_rank_idx) = sf_byte; + } + } else { + // Default BF16 path (16 bytes/lane = 8 BF16). + const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx) + .get_data_buffer(dst_token_idx); + const auto dst_ptr = math::advance_ptr( + dst_token.get_base_ptr(), + n_idx * static_cast(sizeof(nv_bfloat16)) + (lane_idx % 16) * static_cast(sizeof(float4))); + *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + } } } @@ -1590,80 +1688,166 @@ sm100_fp8_fp4_mega_moe_impl(void* y, static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : -1; const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); + // Stream B: FP8 path loads kNumChunkBytes / 2 per slot (FP8 = 1 byte/elem) + // and reads a per-(slot, token, n_block) UE8M0 SF byte to dequant + // on the fly. Output is BF16 either way → store byte count is + // kNumChunkBytes regardless. + constexpr uint32_t kNumLoadBytesPerChunk = + kUseFp8Combine ? (kNumChunkBytes / 2) : kNumChunkBytes; + constexpr uint32_t kNumLoadUint4PerLane = + kUseFp8Combine ? (kNumUint4PerLane / 2) : kNumUint4PerLane; + // Per-uint4 load: BF16 → 8 BF16 = 4 float2 pairs. + // FP8 → 16 FP8 = 8 float2 pairs (dequant'd). + constexpr uint32_t kNumF32PairsPerLoadUint4 = + kUseFp8Combine ? 8u : 4u; + // Per-element offset in the chunk for SF lookup: + // sf_idx = (chunk * kNumLoadElemsPerChunk + elem_in_chunk) / 128 + constexpr uint32_t kNumLoadElemsPerChunk = kHidden / kNumChunks; + // Iterate all chunks for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { - const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; + const uint32_t chunk_byte_offset = chunk * kNumLoadBytesPerChunk; + + // Per-slot SF base pointer cache (FP8 path only; BF16 path leaves these unused). + // We re-read on each slot iteration via __ldg below — values are in L1. + const uint8_t* current_sf_ptr = nullptr; // Move mask and load uint32_t mask = total_mask; - const auto move_mask_and_load = [&](const uint32_t& i) { + const auto move_mask_and_load = [&](const uint32_t& i) -> int { if (mask) { // Move const uint32_t slot_idx = __ffs(mask) - 1; mask ^= 1 << slot_idx; - // Load + // Load FP8 / BF16 chunk if (cute::elect_one_sync()) { const auto src_ptr = math::advance_ptr( combine_token_buffer.get_rank_buffer(slot_idx) .get_data_buffer(token_idx).get_base_ptr(), chunk_byte_offset); - ptx::tma_load_1d(combine_load_buffer[i], src_ptr, combine_load_barriers[i], kNumChunkBytes); - ptx::mbarrier_arrive_and_set_tx(combine_load_barriers[i], kNumChunkBytes); + ptx::tma_load_1d(combine_load_buffer[i], src_ptr, combine_load_barriers[i], kNumLoadBytesPerChunk); + ptx::mbarrier_arrive_and_set_tx(combine_load_barriers[i], kNumLoadBytesPerChunk); } __syncwarp(); - return true; + return static_cast(slot_idx); } - return false; + return -1; }; // Load the first selection - bool do_reduce = move_mask_and_load(load_stage_idx); + int active_slot = move_mask_and_load(load_stage_idx); // Accumulate all top-k contributions for this chunk in float registers float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {}; - while (do_reduce) { - // Prefetch next top-k into the buffer while current is being accumulated - do_reduce = move_mask_and_load(load_stage_idx ^ 1); + while (active_slot >= 0) { + // Prefetch next top-k into the buffer while current is being accumulated. + int next_slot = move_mask_and_load(load_stage_idx ^ 1); - // Accumulate + // Wait for current slot's load. combine_load_barriers[load_stage_idx]->wait(combine_phase); - #pragma unroll - for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { - const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx]; - const auto bf16_values = reinterpret_cast(&uint4_values); + + if constexpr (kUseFp8Combine) { + // Per-slot SF base for this token. + const uint8_t* sf_token_ptr = + combine_sf_buffer.get_rank_buffer(static_cast(active_slot)) + .get_data_buffer(token_idx).get_base_ptr(); #pragma unroll - for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) - ptx::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); + for (uint32_t j = 0; j < kNumLoadUint4PerLane; ++ j) { + const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx]; + // SF for the 16 elements at offset (chunk_elem_offset + (j*32 + lane)*16); + // since 16 < 128, all 16 elements share a single SF byte. + const uint32_t sf_idx = + (chunk * kNumLoadElemsPerChunk + (j * 32 + lane_idx) * 16) / kCombineGranK; + const uint8_t sf_byte = __ldg(sf_token_ptr + sf_idx); + const float sf = math::fast_pow2(static_cast(sf_byte) - 127); + const uint32_t* w = reinterpret_cast(&uint4_values); + #pragma unroll + for (uint32_t l = 0; l < 4; ++ l) { + // Each uint32 = 4 FP8 = 2 FP8x2 — convert via FP16 intermediate. + uint32_t f16_lo, f16_hi; + asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" + : "=r"(f16_lo) : "h"(uint16_t(w[l] & 0xFFFFu))); + asm volatile("cvt.rn.f16x2.e4m3x2 %0, %1;" + : "=r"(f16_hi) : "h"(uint16_t(w[l] >> 16))); + float vlx, vly, vhx, vhy; + asm volatile("cvt.f32.f16 %0, %1;" : "=f"(vlx) : "h"(uint16_t(f16_lo & 0xFFFFu))); + asm volatile("cvt.f32.f16 %0, %1;" : "=f"(vly) : "h"(uint16_t(f16_lo >> 16))); + asm volatile("cvt.f32.f16 %0, %1;" : "=f"(vhx) : "h"(uint16_t(f16_hi & 0xFFFFu))); + asm volatile("cvt.f32.f16 %0, %1;" : "=f"(vhy) : "h"(uint16_t(f16_hi >> 16))); + auto& acc_lo = reduced[j * kNumF32PairsPerLoadUint4 + l * 2 + 0]; + auto& acc_hi = reduced[j * kNumF32PairsPerLoadUint4 + l * 2 + 1]; + acc_lo.x = __fmaf_rn(vlx, sf, acc_lo.x); + acc_lo.y = __fmaf_rn(vly, sf, acc_lo.y); + acc_hi.x = __fmaf_rn(vhx, sf, acc_hi.x); + acc_hi.y = __fmaf_rn(vhy, sf, acc_hi.y); + } + } + } else { + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx]; + const auto bf16_values = reinterpret_cast(&uint4_values); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + ptx::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); + } } combine_phase ^= load_stage_idx; load_stage_idx ^= 1; + active_slot = next_slot; } - // Cast - #pragma unroll - for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { - uint4 casted; - auto casted_bf16 = reinterpret_cast(&casted); + // Cast & write to smem store-buffer. + // BF16 path: kNumUint4PerLane stores, mapping accumulator[j*4+l] → store-uint4 j. + // FP8 path: kNumLoadUint4PerLane * 2 stores, mapping accumulator[j*8+l] → store-uint4 (j*2 + (l/4)). + if constexpr (kUseFp8Combine) { #pragma unroll - for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) - casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); - - // Wait share memory release and write - if (j == 0) { - ptx::tma_store_wait<0>(); - __syncwarp(); + for (uint32_t j = 0; j < kNumLoadUint4PerLane; ++ j) { + // Lower BF16 uint4 (8 elements: pairs 0..3 of this input uint4). + uint4 lo, hi; + auto lo_bf = reinterpret_cast(&lo); + auto hi_bf = reinterpret_cast(&hi); + #pragma unroll + for (uint32_t l = 0; l < 4; ++ l) { + lo_bf[l] = __float22bfloat162_rn(reduced[j * kNumF32PairsPerLoadUint4 + l]); + hi_bf[l] = __float22bfloat162_rn(reduced[j * kNumF32PairsPerLoadUint4 + 4 + l]); + } + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); + } + // Layout: each input uint4 j (16 elements) → 2 BF16 uint4 at + // indices (j * 32 + lane_idx) * 2 + {0, 1}. + ptx::st_shared(combine_store_buffer + (j * 32 + lane_idx) * 2 + 0, + lo.x, lo.y, lo.z, lo.w); + ptx::st_shared(combine_store_buffer + (j * 32 + lane_idx) * 2 + 1, + hi.x, hi.y, hi.z, hi.w); + } + } else { + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + uint4 casted; + auto casted_bf16 = reinterpret_cast(&casted); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); + } + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, + casted.x, casted.y, casted.z, casted.w); } - ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, - casted.x, casted.y, casted.z, casted.w); } __syncwarp(); - // TMA store the token chunk + // TMA store the BF16 chunk to gmem y. Output byte offset still + // tracks BF16 chunks (= kNumChunkBytes). if (cute::elect_one_sync()) { cute::tma_store_fence(); ptx::tma_store_1d( - math::advance_ptr(y, static_cast(token_idx) * kNumHiddenBytes + chunk_byte_offset), + math::advance_ptr(y, static_cast(token_idx) * kNumHiddenBytes + chunk * kNumChunkBytes), combine_store_buffer, kNumChunkBytes); cute::tma_store_arrive(); } From 8ab00962e977c2fb0d5f8c7f7366277d7b2aee70 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Tue, 12 May 2026 15:51:14 -0700 Subject: [PATCH 09/26] Add tvm-ffi wrapper for w4a4 megamoe --- csrc/tvm_ffi_api.cpp | 21 +++++++++++++++++++++ csrc/utils/torch_compat.hpp | 14 ++++++++++++++ sgl_deep_gemm/__init__.py | 1 + 3 files changed, 36 insertions(+) diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index b992c05b8d..a029070792 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -579,9 +579,30 @@ void dg_fp8_fp4_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weig ); } + +void dg_mega_moe_pre_dispatch( + TensorView x, TensorView topk_idx, TensorView topk_weights, + TensorView buf_x, TensorView buf_x_sf, + TensorView buf_topk_idx, TensorView buf_topk_weights, + int64_t num_tokens, int64_t group_size, bool use_fp4_acts) { + mega_moe_pre_dispatch( + convert_to_torch_tensor(x), + convert_to_torch_tensor(topk_idx), + convert_to_torch_tensor(topk_weights), + convert_to_torch_tensor(buf_x), + convert_to_torch_tensor(buf_x_sf), + convert_to_torch_tensor(buf_topk_idx), + convert_to_torch_tensor(buf_topk_weights), + static_cast(num_tokens), + static_cast(group_size), + use_fp4_acts + ); +} + TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_token_alignment_for_mega_moe, dg_get_token_alignment_for_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_symm_buffer_size_for_mega_moe, dg_get_symm_buffer_size_for_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_mega_moe, dg_fp8_fp4_mega_moe); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(mega_moe_pre_dispatch, dg_mega_moe_pre_dispatch); #endif // DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE diff --git a/csrc/utils/torch_compat.hpp b/csrc/utils/torch_compat.hpp index 14fd399d17..c9199ede95 100644 --- a/csrc/utils/torch_compat.hpp +++ b/csrc/utils/torch_compat.hpp @@ -131,6 +131,20 @@ inline torch::Tensor convert_to_torch_tensor(tvm::ffi::TensorView tensor) { .device(torch::kCUDA, device_id) .requires_grad(false); + // Zero-numel tensors may carry a nullptr data_ptr, which trips + // torch::from_blob() getDeviceFromPtr() host-vs-device check. + // Allocate a fresh empty CUDA tensor in that case: any kernel reading + // through it must already guard on the zero-size dimension. + int64_t numel = 1; + for (auto s : sizes) numel *= s; + if (numel == 0 || data == nullptr) { + if (tensor.strides().data()) { + auto strides = std::vector(tensor.strides().begin(), tensor.strides().end()); + return torch::empty_strided(sizes, strides, opts); + } + return torch::empty(sizes, opts); + } + if (tensor.strides().data()) { auto strides = std::vector(tensor.strides().begin(), tensor.strides().end()); return torch::from_blob(data, sizes, strides, opts); diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index 833d9ae67d..0a3f8dbe37 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -262,6 +262,7 @@ def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe, fp8_fp4_mega_moe, + mega_moe_pre_dispatch, ) # Some utils From 89f2f0038df3030369ea5fdf154b03af1ee9bbbb Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Tue, 12 May 2026 16:14:32 -0700 Subject: [PATCH 10/26] Bump to v0.1.0 --- sgl_deep_gemm/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sgl_deep_gemm/VERSION b/sgl_deep_gemm/VERSION index 16acc60685..6c6aa7cb09 100644 --- a/sgl_deep_gemm/VERSION +++ b/sgl_deep_gemm/VERSION @@ -1 +1 @@ -0.0.1rc1 \ No newline at end of file +0.1.0 \ No newline at end of file From 23105b21d4b52470e1d44a5ece2b033c10bb2b17 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Thu, 14 May 2026 00:38:41 -0700 Subject: [PATCH 11/26] update readme --- sgl_deep_gemm/README.md | 18 ++++++++++++++++-- 1 file changed, 16 insertions(+), 2 deletions(-) diff --git a/sgl_deep_gemm/README.md b/sgl_deep_gemm/README.md index 64064fe41e..39d56d391e 100644 --- a/sgl_deep_gemm/README.md +++ b/sgl_deep_gemm/README.md @@ -1,2 +1,16 @@ -sgl-deep-gemm is a package built from SGLang's customized branch of DeepGemm. -To build it locally, run `bash build_sgl_deep_gemm.sh`, then pip install the wheel generated under `dist`. \ No newline at end of file + +## Introduction + +sgl-deep-gemm is a pypi package built from SGLang's customized branch of DeepGemm. Comparing with origina DeepGemm, it supports the following features to better support SGLang: +1. ABI support: with the help of tvm-ffi wrappers, a single wheel can run on different python versions. +2. pypi support: easy installation with `pip install sgl-deep-gemm`. No need to manually search for wheel links. +3. Fast iteration: add custom kernels and bump versions at no time. + +## Usage +To build it locally, run `bash build_sgl_deep_gemm.sh`, then pip install the wheel generated under `dist`. + +To release a new set of wheels, please contact SGLang team and run the [release workflow](https://github.com/sgl-project/sglang/actions/workflows/release-whl-deepgemm.yml) under SGLang repo + +For each major version release (0.X.Y -> 0.(X+1).0), a new branch should be created (release/v0.(X+1).0) for stability purpose. + +For any incoming pull requests, it should be rebased upon `dev` branch. From 46a229498f9772f18966d7bc77b3374b8cdd40e1 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Wed, 20 May 2026 02:56:55 -0400 Subject: [PATCH 12/26] Export the PDL utils of DeepGEMM (#34) Co-authored-by: b8zhong --- csrc/tvm_ffi_api.cpp | 4 ++++ deep_gemm/__init__.py | 2 ++ 2 files changed, 6 insertions(+) diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index a029070792..d82e382e29 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -41,6 +41,8 @@ void dg_set_num_sms(int64_t n) { device_runtime->set_num_sms(static_cast(n) // void dg_set_compile_mode(int64_t n) { device_runtime->set_compile_mode(static_cast(n)); } int64_t dg_get_tc_util() { return device_runtime->get_tc_util(); } void dg_set_tc_util(int64_t n) { device_runtime->set_tc_util(static_cast(n)); } +bool dg_get_pdl() { return device_runtime->get_pdl(); } +void dg_set_pdl(bool v) { device_runtime->set_pdl(v); } TVM_FFI_DLL_EXPORT_TYPED_FUNC(init, dg_init); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_num_sms, dg_get_num_sms); @@ -49,6 +51,8 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_num_sms, dg_set_num_sms); // TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_compile_mode, dg_set_compile_mode); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_tc_util, dg_get_tc_util); TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_tc_util, dg_set_tc_util); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_pdl, dg_get_pdl); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(set_pdl, dg_set_pdl); // --------------------------------------------------------------------------- // Layout utilities diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 6c0f6a769c..66ee039ef4 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -143,6 +143,8 @@ def _load_module() -> Module: # get_compile_mode = _C.get_compile_mode set_tc_util = _C.set_tc_util get_tc_util = _C.get_tc_util +set_pdl = _C.set_pdl +get_pdl = _C.get_pdl # cuBLASLt Kernels cublaslt_gemm_nt = _C.cublaslt_gemm_nt From 48eb6a67ca99feaabbd7fd1d4a7f4f8b3ac3e381 Mon Sep 17 00:00:00 2001 From: popsiclexu Date: Thu, 28 May 2026 11:19:37 +0800 Subject: [PATCH 13/26] Expose BF16 grouped GEMM wrappers --- csrc/tvm_ffi_api.cpp | 26 ++++++++++++++++++++++++++ deep_gemm/__init__.py | 9 ++++++++- sgl_deep_gemm/__init__.py | 9 ++++++++- 3 files changed, 42 insertions(+), 2 deletions(-) diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index d82e382e29..78649118fb 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -386,6 +386,30 @@ void dg_bf16_gemm_tt(TensorView a, TensorView b, TensorView d, gemm::bf16_gemm_tt(convert_to_torch_tensor(a), convert_to_torch_tensor(b), convert_to_torch_tensor(d), c_opt, compiled_dims); } +void dg_m_grouped_bf16_gemm_nt_contiguous(TensorView a, TensorView b, TensorView d, + TensorView grouped_layout, + std::string compiled_dims, + bool use_psum_layout, + Optional expected_m_for_psum_layout) { + auto expected_m_opt = expected_m_for_psum_layout.has_value()? std::make_optional((int)expected_m_for_psum_layout.value()) : std::nullopt; + gemm::m_grouped_bf16_gemm_nt_contiguous( + convert_to_torch_tensor(a), convert_to_torch_tensor(b), + convert_to_torch_tensor(d), convert_to_torch_tensor(grouped_layout), + compiled_dims, use_psum_layout, expected_m_opt + ); +} + +void dg_m_grouped_bf16_gemm_nt_masked(TensorView a, TensorView b, TensorView d, + TensorView masked_m, + int64_t expected_m, + std::string compiled_dims) { + gemm::m_grouped_bf16_gemm_nt_masked( + convert_to_torch_tensor(a), convert_to_torch_tensor(b), + convert_to_torch_tensor(d), convert_to_torch_tensor(masked_m), + (int) expected_m, compiled_dims + ); +} + TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_nt, dg_fp8_fp4_gemm_nt); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_nn, dg_fp8_fp4_gemm_nn); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_tn, dg_fp8_fp4_gemm_tn); @@ -397,6 +421,8 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_nt, dg_bf16_gemm_nt); TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_nn, dg_bf16_gemm_nn); TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_tn, dg_bf16_gemm_tn); TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_tt, dg_bf16_gemm_tt); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_bf16_gemm_nt_contiguous, dg_m_grouped_bf16_gemm_nt_contiguous); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_bf16_gemm_nt_masked, dg_m_grouped_bf16_gemm_nt_masked); // Einsum void dg_einsum(std::string expr, TensorView a, TensorView b, TensorView d, diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 66ee039ef4..2bce4038da 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -254,7 +254,14 @@ def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_fp4_gemm_nt_masked - bf16_m_grouped_gemm_nt_masked = None + + def m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, compiled_dims='nk', use_psum_layout=False, expected_m_for_psum_layout=None): + _C.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, compiled_dims, use_psum_layout, expected_m_for_psum_layout) + + def m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m, compiled_dims='nk'): + _C.m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m, compiled_dims) + + bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked except AttributeError: pass diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index 0a3f8dbe37..50750112f7 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -251,7 +251,14 @@ def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_fp4_gemm_nt_masked - bf16_m_grouped_gemm_nt_masked = None + + def m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, compiled_dims='nk', use_psum_layout=False, expected_m_for_psum_layout=None): + _C.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, compiled_dims, use_psum_layout, expected_m_for_psum_layout) + + def m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m, compiled_dims='nk'): + _C.m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m, compiled_dims) + + bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked except AttributeError: pass From 6c9eaca07a8799fb364a855c0c4b1af2715c078d Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Fri, 29 May 2026 00:55:41 -0700 Subject: [PATCH 14/26] Fix version of dev branch to 0.0.0 --- sgl_deep_gemm/VERSION | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sgl_deep_gemm/VERSION b/sgl_deep_gemm/VERSION index 6c6aa7cb09..bd52db81d0 100644 --- a/sgl_deep_gemm/VERSION +++ b/sgl_deep_gemm/VERSION @@ -1 +1 @@ -0.1.0 \ No newline at end of file +0.0.0 \ No newline at end of file From 86d705dfdbc0d461b71d908fb4390b0a033229d1 Mon Sep 17 00:00:00 2001 From: nvjullin Date: Tue, 2 Jun 2026 04:37:24 +0800 Subject: [PATCH 15/26] Fix IMA guard in paged MQA logits scheduler (#38) Signed-off-by: Jin Li <59594262+liji-nv@users.noreply.github.com> Co-authored-by: Jin Li <59594262+liji-nv@users.noreply.github.com> --- deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh index 548bbbc6ba..8957fdc6a1 100644 --- a/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/paged_mqa_logits.cuh @@ -187,7 +187,7 @@ struct PagedMQALogitsScheduler : IndicesStorage { current_q_atom_idx = current_pack.x, current_kv_idx = current_pack.y * kNumBlocksPerSplit; end_q_atom_idx = end_pack.x, end_kv_idx = end_pack.y * kNumBlocksPerSplit; - current_num_kv = get_num_kv(current_q_atom_idx); + current_num_kv = exist_q_atom_idx(current_q_atom_idx) ? get_num_kv(current_q_atom_idx) : 0; } // Advance step in q_atom_idx space when moving to the next atom. From 09fc8106286ef6eccf58c3648b4c4de7c9893c25 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 8 Jun 2026 18:48:17 -0700 Subject: [PATCH 16/26] Fix various issues in DeepGEMM tests (#39) Co-authored-by: Brayden Zhong Co-authored-by: Claude Opus 4.8 (1M context) --- csrc/tvm_ffi_api.cpp | 29 ++ deep_gemm/mega/__init__.py | 5 +- sgl_deep_gemm/__init__.py | 6 + sgl_deep_gemm/tests/generators.py | 407 ++++++++++++++++++ sgl_deep_gemm/tests/test_attention.py | 397 +++++++++++++++++ sgl_deep_gemm/tests/test_bf16.py | 223 ++++++++++ sgl_deep_gemm/tests/test_einsum.py | 181 ++++++++ sgl_deep_gemm/tests/test_fp8_fp4.py | 223 ++++++++++ sgl_deep_gemm/tests/test_hyperconnection.py | 57 +++ sgl_deep_gemm/tests/test_layout.py | 112 +++++ sgl_deep_gemm/tests/test_lazy_init.py | 20 + sgl_deep_gemm/tests/test_legacy.py | 90 ++++ sgl_deep_gemm/tests/test_mega_moe.py | 301 +++++++++++++ .../tests}/test_mega_moe_l1_fp4_accuracy.py | 51 ++- .../tests}/test_mega_moe_l1_sentinel.py | 0 .../tests}/test_mega_moe_pre_dispatch.py | 0 sgl_deep_gemm/tests/test_sanitizer.py | 79 ++++ tests/generators.py | 3 + tests/test_attention.py | 51 ++- tests/test_bf16.py | 10 +- tests/test_fp8_fp4.py | 11 +- tests/test_mega_moe.py | 139 +++--- 22 files changed, 2275 insertions(+), 120 deletions(-) create mode 100644 sgl_deep_gemm/tests/generators.py create mode 100644 sgl_deep_gemm/tests/test_attention.py create mode 100644 sgl_deep_gemm/tests/test_bf16.py create mode 100644 sgl_deep_gemm/tests/test_einsum.py create mode 100644 sgl_deep_gemm/tests/test_fp8_fp4.py create mode 100644 sgl_deep_gemm/tests/test_hyperconnection.py create mode 100644 sgl_deep_gemm/tests/test_layout.py create mode 100644 sgl_deep_gemm/tests/test_lazy_init.py create mode 100644 sgl_deep_gemm/tests/test_legacy.py create mode 100644 sgl_deep_gemm/tests/test_mega_moe.py rename {tests => sgl_deep_gemm/tests}/test_mega_moe_l1_fp4_accuracy.py (95%) rename {tests => sgl_deep_gemm/tests}/test_mega_moe_l1_sentinel.py (100%) rename {tests => sgl_deep_gemm/tests}/test_mega_moe_pre_dispatch.py (100%) create mode 100644 sgl_deep_gemm/tests/test_sanitizer.py diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index 78649118fb..abd46179cd 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -399,6 +399,17 @@ void dg_m_grouped_bf16_gemm_nt_contiguous(TensorView a, TensorView b, TensorView ); } +void dg_m_grouped_bf16_gemm_nn_contiguous(TensorView a, TensorView b, TensorView d, + TensorView grouped_layout, + std::string compiled_dims, + bool use_psum_layout) { + gemm::m_grouped_bf16_gemm_nn_contiguous( + convert_to_torch_tensor(a), convert_to_torch_tensor(b), + convert_to_torch_tensor(d), convert_to_torch_tensor(grouped_layout), + compiled_dims, use_psum_layout + ); +} + void dg_m_grouped_bf16_gemm_nt_masked(TensorView a, TensorView b, TensorView d, TensorView masked_m, int64_t expected_m, @@ -410,6 +421,22 @@ void dg_m_grouped_bf16_gemm_nt_masked(TensorView a, TensorView b, TensorView d, ); } +void dg_k_grouped_bf16_gemm_tn_contiguous(TensorView a, TensorView b, TensorView d, + Array ks, TensorView ks_tensor, + Optional c, + std::string compiled_dims) { + std::vector ks_val; + ks_val.reserve(ks.size()); + for (Array::iterator it = ks.begin(); it != ks.end(); ++it) + ks_val.push_back(static_cast(*it)); + auto c_opt = c.has_value()? std::optional(convert_to_torch_tensor(c.value())) : std::nullopt; + gemm::k_grouped_bf16_gemm_tn_contiguous( + convert_to_torch_tensor(a), convert_to_torch_tensor(b), + convert_to_torch_tensor(d), ks_val, convert_to_torch_tensor(ks_tensor), + c_opt, compiled_dims + ); +} + TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_nt, dg_fp8_fp4_gemm_nt); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_nn, dg_fp8_fp4_gemm_nn); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_gemm_tn, dg_fp8_fp4_gemm_tn); @@ -422,7 +449,9 @@ TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_nn, dg_bf16_gemm_nn); TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_tn, dg_bf16_gemm_tn); TVM_FFI_DLL_EXPORT_TYPED_FUNC(bf16_gemm_tt, dg_bf16_gemm_tt); TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_bf16_gemm_nt_contiguous, dg_m_grouped_bf16_gemm_nt_contiguous); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_bf16_gemm_nn_contiguous, dg_m_grouped_bf16_gemm_nn_contiguous); TVM_FFI_DLL_EXPORT_TYPED_FUNC(m_grouped_bf16_gemm_nt_masked, dg_m_grouped_bf16_gemm_nt_masked); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(k_grouped_bf16_gemm_tn_contiguous, dg_k_grouped_bf16_gemm_tn_contiguous); // Einsum void dg_einsum(std::string expr, TensorView a, TensorView b, TensorView d, diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 77db53cd92..2b58232d7e 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -41,11 +41,12 @@ def __init__(self, group: dist.ProcessGroup, self.group.barrier() torch.cuda.synchronize() - # Create input buffer views + # Create input buffer views (as torch tensors, not tvm-ffi tensors). (self.x, self.x_sf, self.topk_idx, self.topk_weights, self.l1_acts, self.l1_acts_sf, - self.l2_acts, self.l2_acts_sf) = slice_input_buffers(self.buffer) + self.l2_acts, self.l2_acts_sf) = map( + torch.from_dlpack, slice_input_buffers(self.buffer)) def destroy(self): self.handle = None diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index 50750112f7..12cbf9ffea 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -255,9 +255,15 @@ def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, def m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, compiled_dims='nk', use_psum_layout=False, expected_m_for_psum_layout=None): _C.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, compiled_dims, use_psum_layout, expected_m_for_psum_layout) + def m_grouped_bf16_gemm_nn_contiguous(a, b, d, grouped_layout, compiled_dims='nk', use_psum_layout=False): + _C.m_grouped_bf16_gemm_nn_contiguous(a, b, d, grouped_layout, compiled_dims, use_psum_layout) + def m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m, compiled_dims='nk'): _C.m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m, compiled_dims) + def k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, ks_tensor, c=None, compiled_dims='mn'): + _C.k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, ks_tensor, c, compiled_dims) + bf16_m_grouped_gemm_nt_masked = m_grouped_bf16_gemm_nt_masked except AttributeError: diff --git a/sgl_deep_gemm/tests/generators.py b/sgl_deep_gemm/tests/generators.py new file mode 100644 index 0000000000..989e984e7c --- /dev/null +++ b/sgl_deep_gemm/tests/generators.py @@ -0,0 +1,407 @@ +import enum +import random +import torch +from typing import Generator, List, Optional, Tuple + +from deep_gemm.testing import get_arch_major +from deep_gemm.utils import ( + align, ceil_div, + per_token_cast_to_fp8, per_channel_cast_to_fp8, per_block_cast_to_fp8, + per_token_cast_to_fp4, transpose_packed_fp4, + get_mk_alignment_for_contiguous_layout, + set_mk_alignment_for_contiguous_layout +) + + +class KernelType(enum.Enum): + Kernel1D1D = 0 + Kernel1D2D = 1 + KernelNoSF = 2 + + def is_1d1d(self): + return self.value == 0 + + def is_1d2d(self): + return self.value == 1 + + def is_nosf(self): + return self.value == 2 + + +class MajorTypeAB(enum.Enum): + KMajor = 0 + MNMajor = 1 + + def is_k_major(self): + return self.value == 0 + + def is_mn_major(self): + return self.value == 1 + + +class QuantConfig: + _legacy_quant_config = (128, 128, False, False) + + def __init__(self, value: Tuple[int, int, bool, bool] = _legacy_quant_config): + self.gran_k_a, self.gran_k_b, self.is_fp4_a, self.is_fp4_b = value + + def print(self): + print(f' > Testing with gran_k_a={self.gran_k_a}, gran_k_b={self.gran_k_b}, ' + f'is_fp4_a={self.is_fp4_a}, is_fp4_b={self.is_fp4_b}') + + def is_legacy(self) -> bool: + return (self.gran_k_a, self.gran_k_b, self.is_fp4_a, self.is_fp4_b) == self._legacy_quant_config + + def get_recipes(self, is_wgrad: bool = False) -> Tuple[Tuple, Tuple, Tuple]: + recipe, recipe_a, recipe_b = None, None, None + if self.is_legacy(): + recipe = (1, 1, 128) if is_wgrad else None + else: + recipe_a = (1, self.gran_k_a) + recipe_b = (1, self.gran_k_b) if self.is_fp4_b or is_wgrad else (self.gran_k_b, self.gran_k_b) + return recipe, recipe_a, recipe_b + + def max_diff(self) -> float: + if self.is_fp4_a and self.is_fp4_b: + return 0.02 + if self.is_fp4_a or self.is_fp4_b: + return 0.01 + return 0.001 + + @staticmethod + def get_list_from_dtype(dtype: torch.dtype) -> List: + if dtype == torch.bfloat16: + return [None] + quant_config_list = [QuantConfig()] + if get_arch_major() == 10: + quant_config_list.append(QuantConfig((128, 32, False, True))) + return quant_config_list + + +def reset_seed(seed: int = 0): + random.seed(seed) + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + + +def get_ue8m0_usage(kernel_type: KernelType) -> bool: + if get_arch_major() == 9: + return False + return kernel_type.is_1d1d() + + +def get_kernel_types(dtype: torch.dtype) -> tuple: + if dtype == torch.bfloat16: + return (KernelType.KernelNoSF, ) + + return (KernelType.Kernel1D2D, ) if get_arch_major() == 9 else (KernelType.Kernel1D1D, ) + + +def get_major_ab(allow_a_mn_major: bool, allow_b_mn_major: bool) -> Generator: + for major_a in (MajorTypeAB.KMajor, MajorTypeAB.MNMajor): + for major_b in (MajorTypeAB.KMajor, MajorTypeAB.MNMajor): + if major_a.is_mn_major() and not allow_a_mn_major: + continue + if major_b.is_mn_major() and not allow_b_mn_major: + continue + yield major_a, major_b + + +def get_psum_layout_usage() -> tuple: + return True, False + + +def enumerate_normal(dtype: torch.dtype) -> Generator: + assert dtype in (torch.float8_e4m3fn, torch.bfloat16) + + quant_config_list = QuantConfig.get_list_from_dtype(dtype) + fp32_output_nk = [(256, 7168), (129280, 7168)] + bf16_output_nk = [(2112, 7168), (576, 7168), (24576, 1536), (32768, 512), (7168, 16384), (4096, 7168), (7168, 2048)] + m_fwd_list, m_bwd_list = [1, 128, 4096], [4096, ] + nk_list = list(bf16_output_nk) + + # Only BF16 GEMM needs FP32 outputs + if dtype == torch.bfloat16: + nk_list += fp32_output_nk + + for kernel_type in get_kernel_types(dtype): + for quant_config in quant_config_list: + if len(quant_config_list) > 1: + quant_config.print() + reset_seed() + + # Forward + for m in m_fwd_list: + for i in range(len(nk_list)): + n, k = nk_list[i] + out_dtype = torch.bfloat16 if i < len(bf16_output_nk) else torch.float + yield kernel_type, quant_config, m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, False, out_dtype + + # Backward + for m in m_bwd_list: + for n, k in nk_list: + override_major = MajorTypeAB.MNMajor + override_kernel_type = kernel_type + if get_arch_major() == 9 and dtype == torch.float8_e4m3fn: + override_major = MajorTypeAB.KMajor + override_kernel_type = KernelType.Kernel1D1D + yield kernel_type, quant_config, m, k, n, MajorTypeAB.KMajor, override_major, False, torch.bfloat16 # Dgrad + yield override_kernel_type, quant_config, n, m, k, override_major, override_major, True, torch.float # Wgrad + yield override_kernel_type, quant_config, n, m, k, override_major, override_major, False, torch.bfloat16 # Wgrad + + +def enumerate_m_grouped_contiguous(dtype: torch.dtype) -> Generator: + quant_config_list = QuantConfig.get_list_from_dtype(dtype) + m_group_list = [(4, 8192), (8, 4096)] + n_k_list = [(6144, 7168), (7168, 3072), (4096, 4096), (4096, 2048)] + for kernel_type in get_kernel_types(dtype): + for quant_config in quant_config_list: + if len(quant_config_list) > 1: + quant_config.print() + for use_psum_layout in get_psum_layout_usage(): + reset_seed() + for num_groups, expected_m_per_group in m_group_list: + for n, k in n_k_list: + for major_a, major_b in get_major_ab(False, get_arch_major() != 9 or dtype != torch.float8_e4m3fn): + yield kernel_type, quant_config, num_groups, expected_m_per_group, n, k, major_a, major_b, use_psum_layout + + +def enumerate_m_grouped_masked(dtype: torch.dtype) -> Generator: + quant_config_list = QuantConfig.get_list_from_dtype(dtype) + max_m = 4096 + m_group_list = [(32, 192), (6, 1024), (32, 20), (6, 20)] + n_k_list = [(6144, 7168), (7168, 3072), (4096, 4096), (4096, 2048)] + for kernel_type in get_kernel_types(dtype): + for quant_config in quant_config_list: + if len(quant_config_list) > 1: + quant_config.print() + for use_psum_layout in get_psum_layout_usage(): + reset_seed() + for num_groups, m in m_group_list: + for n, k in n_k_list: + yield kernel_type, quant_config, num_groups, max_m, m, n, k, use_psum_layout + + +def enumerate_k_grouped_contiguous(dtype: torch.dtype): + gran_k_list = (128, ) if get_arch_major() == 9 else (32, 128) + # Only K-major is supported for SM90 FP8 + major_a, major_b = (MajorTypeAB.KMajor, MajorTypeAB.KMajor) if get_arch_major() == 9 and dtype == torch.float8_e4m3fn \ + else (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor) + # Must with FP32 accumulation and 1D1D kernels + for num_groups, m, n, expected_k_per_group in (( 4, 4096, 7168, 8192), ( 4, 7168, 2048, 8192), # EP64 + ( 8, 4096, 7168, 4096), ( 8, 7168, 2048, 4096), # EP32 + (16, 4096, 7168, 2048), (16, 7168, 2048, 2048)): # EP16 + if dtype == torch.bfloat16: + ks = [align(int(expected_k_per_group * random.uniform(0.7, 1.3)), get_mk_alignment_for_contiguous_layout()) for _ in range(num_groups)] + yield num_groups, m, n, major_a, major_b, ks, expected_k_per_group + else: + for gran_k in gran_k_list: + set_mk_alignment_for_contiguous_layout(gran_k) + ks = [align(int(expected_k_per_group * random.uniform(0.7, 1.3)), gran_k) for _ in range(num_groups)] + yield num_groups, m, n, major_a, major_b, ks, expected_k_per_group, gran_k + + +def enumerate_sf_layout(): + gran_k_list = (128, ) if get_arch_major() == 9 else (32, 128) + for use_ue8m0 in (False, True): + for with_transpose in (True, False): + for mn in (4096, 4097, 8192): + for k in (128, 7168, 7296): + for num_groups in (1, 2, 4): + for gran_k in gran_k_list: + set_mk_alignment_for_contiguous_layout(gran_k) + yield mn, k, with_transpose, use_ue8m0, num_groups, gran_k + + +def enumerate_k_grouped_sf_layout(): + gran_k_list = (128, ) if get_arch_major() == 9 else (32, 128) + for mn in (4096, 7168): + for num_groups, avg_k in ((16, 2048), (8, 4096), (72, 384), (128, 256)): + for gran_k in gran_k_list: + set_mk_alignment_for_contiguous_layout(gran_k) + ks = [align(int(random.uniform(0.7, 1.3) * avg_k), gran_k) for _ in range(num_groups)] + yield mn, ks, num_groups, gran_k + + +def enumerate_transpose(): + for mn in (64, 4096, 16384): + for delta in (0, 101, 202, 303): + for k in (128, 1024, 4096, 9984, 16384): + yield mn + delta, k + + +def cast_fp8_fp4_with_major(x: torch.Tensor, major: MajorTypeAB, gran_k: int, is_fp4: bool, + use_ue8m0: bool, use_block_cast_for_fp8: bool = False): + if is_fp4: + x_fp4 = per_token_cast_to_fp4(x, use_ue8m0=use_ue8m0, gran_k=gran_k) + return x_fp4 if major.is_k_major() else (transpose_packed_fp4(x_fp4[0]).T, x_fp4[1]) + else: + x_fp8 = per_block_cast_to_fp8(x, use_ue8m0=use_ue8m0, gran_k=gran_k) if use_block_cast_for_fp8 \ + else per_token_cast_to_fp8(x, use_ue8m0=use_ue8m0, gran_k=gran_k) + return x_fp8 if major.is_k_major() else (x_fp8[0].T.contiguous().T, x_fp8[1]) + + +def grouped_cast_fp8_fp4_with_major(x: torch.Tensor, major: MajorTypeAB, gran_k: int, is_fp4: bool, + use_ue8m0: bool, use_block_cast_for_fp8: bool = False): + num_groups, mn, k = x.size() + if is_fp4: + x_fp4 = (torch.empty((num_groups, mn, k // 2), device='cuda', dtype=torch.int8) if major.is_k_major() else \ + torch.empty((num_groups, k, mn // 2), device='cuda', dtype=torch.int8), + torch.empty((num_groups, mn, ceil_div(k, gran_k)), device='cuda', dtype=torch.float)) + for i in range(num_groups): + x_i_fp4 = per_token_cast_to_fp4(x[i], use_ue8m0=use_ue8m0, gran_k=gran_k) + x_fp4[0][i], x_fp4[1][i] = x_i_fp4 if major.is_k_major() else (transpose_packed_fp4(x_i_fp4[0]), x_i_fp4[1]) + return x_fp4 if major.is_k_major() else (x_fp4[0].mT, x_fp4[1]) + else: + x_fp8 = (torch.empty_like(x, dtype=torch.float8_e4m3fn), + torch.empty((num_groups, ceil_div(mn, gran_k), ceil_div(k, gran_k)), device='cuda', dtype=torch.float) if use_block_cast_for_fp8 \ + else torch.empty((num_groups, mn, ceil_div(k, gran_k)), device='cuda', dtype=torch.float)) + for i in range(num_groups): + x_fp8[0][i], x_fp8[1][i] = per_block_cast_to_fp8(x[i], use_ue8m0=use_ue8m0, gran_k=gran_k) if use_block_cast_for_fp8 \ + else per_token_cast_to_fp8(x[i], use_ue8m0=use_ue8m0, gran_k=gran_k) + return x_fp8 if major.is_k_major() else (x_fp8[0].mT.contiguous().mT, x_fp8[1]) + + +def generate_normal(m: int, n: int, k: int, + major_a: MajorTypeAB, major_b: MajorTypeAB, + accumulate: bool, out_dtype: torch.dtype, + kernel_type: KernelType, + use_ue8m0: bool = False, use_bf16: bool = False, + quant_config: Optional[QuantConfig] = None): + a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) + b = torch.randn((n, k), device='cuda', dtype=torch.bfloat16) + d = torch.randn((m, n), device='cuda', dtype=out_dtype) * 32 if accumulate else \ + torch.empty((m, n), device='cuda', dtype=out_dtype) + c = d if accumulate else None + ref_d = (a.float() @ b.float().t() + (c if accumulate else 0)).to(out_dtype) + + if use_bf16: + a = a if major_a.is_k_major() else a.T.contiguous().T + b = b if major_b.is_k_major() else b.T.contiguous().T + return a, b, c, d, ref_d + + quant_config = QuantConfig() if quant_config is None else quant_config + a = cast_fp8_fp4_with_major(a, major_a, quant_config.gran_k_a, quant_config.is_fp4_a, use_ue8m0) + b = cast_fp8_fp4_with_major(b, major_b, quant_config.gran_k_b, quant_config.is_fp4_b, use_ue8m0, + use_block_cast_for_fp8=not (kernel_type.is_1d1d() and accumulate)) + + return a, b, c, d, ref_d + + +def generate_m_grouped_contiguous(num_groups: int, expected_m_per_group: int, n: int, k: int, + major_a: MajorTypeAB, major_b: MajorTypeAB, + use_ue8m0: bool = False, use_bf16: bool = False, + use_psum_layout: bool = False, + quant_config: Optional[QuantConfig] = None): + actual_ms = [int(expected_m_per_group * random.uniform(0.7, 1.3)) for _ in range(num_groups)] + aligned_ms = [align(actual_m, get_mk_alignment_for_contiguous_layout()) for actual_m in actual_ms] + m = sum(aligned_ms) + + a = torch.randn((m, k), device='cuda', dtype=torch.bfloat16) + b = torch.randn((num_groups, n, k), device='cuda', dtype=torch.bfloat16) + grouped_layout = torch.empty(num_groups, device='cuda', dtype=torch.int32) if use_psum_layout \ + else torch.empty(m, device='cuda', dtype=torch.int32) + d = torch.empty((m, n), device='cuda', dtype=torch.bfloat16) + ref_d = torch.randn((m, n), device='cuda', dtype=torch.bfloat16) + + start = 0 + for i, (actual_m, aligned_m) in enumerate(zip(actual_ms, aligned_ms)): + actual_end = start + actual_m + aligned_end = start + aligned_m + if use_psum_layout: + grouped_layout[i] = actual_end + else: + grouped_layout[start: actual_end] = i + grouped_layout[actual_end: aligned_end] = -1 + a[actual_end: aligned_end] = 0 + ref_d[start: aligned_end] = a[start: aligned_end] @ b[i].t() + start = aligned_end + + if use_bf16: + b = b if major_b.is_k_major() else b.mT.contiguous().mT + return m, a, b, grouped_layout, d, ref_d + + assert major_a.is_k_major() + quant_config = QuantConfig() if quant_config is None else quant_config + a = cast_fp8_fp4_with_major(a, major_a, quant_config.gran_k_a, quant_config.is_fp4_a, use_ue8m0) + b = grouped_cast_fp8_fp4_with_major(b, major_b, quant_config.gran_k_b, quant_config.is_fp4_b, use_ue8m0, use_block_cast_for_fp8=True) + + return m, a, b, grouped_layout, d, ref_d + + +def layout_masked_to_psum(x: torch.Tensor, psum_m: torch.Tensor): + num_groups, max_m, _ = x.size() + x_psum = torch.empty_like(x).view(num_groups * max_m, -1) + last_psum_m = 0 + for i in range(num_groups): + x_psum[last_psum_m: psum_m[i]] = x[i, :psum_m[i] - last_psum_m] + last_psum_m = align(psum_m[i], get_mk_alignment_for_contiguous_layout()) + return x_psum + + +def generate_m_grouped_masked(num_groups: int, max_m: int, expected_m_per_group: int, n: int, k: int, + use_ue8m0: bool = False, use_bf16: bool = False, + use_psum_layout: bool = False, + quant_config: Optional[QuantConfig] = None): + a = torch.randn((num_groups, max_m, k), device='cuda', dtype=torch.bfloat16) + b = torch.randn((num_groups, n, k), device='cuda', dtype=torch.bfloat16) + d = torch.empty((num_groups, max_m, n), device='cuda', dtype=torch.bfloat16) + ref_d = torch.einsum('gmk,gnk->gmn', a, b) + + masked_m = torch.empty((num_groups, ), device='cuda', dtype=torch.int) + psum_m = torch.empty((num_groups, ), device='cuda', dtype=torch.int) + for j in range(num_groups): + masked_m[j] = int(expected_m_per_group * random.uniform(0.7, 1.3)) + psum_m[j] = (0 if j == 0 else align(psum_m[j - 1], get_mk_alignment_for_contiguous_layout())) + masked_m[j] + assert masked_m.amax().item() <= max_m + + if use_bf16: + return a, b, masked_m, psum_m, d, ref_d + + quant_config = QuantConfig() if quant_config is None else quant_config + a = grouped_cast_fp8_fp4_with_major(a, MajorTypeAB.KMajor, quant_config.gran_k_a, quant_config.is_fp4_a, use_ue8m0) + b = grouped_cast_fp8_fp4_with_major(b, MajorTypeAB.KMajor, quant_config.gran_k_b, quant_config.is_fp4_b, use_ue8m0, use_block_cast_for_fp8=True) + + return a, b, masked_m, psum_m, d, ref_d + + +def generate_k_grouped_contiguous(num_groups: int, m: int, n: int, major_a: MajorTypeAB, major_b: MajorTypeAB, ks: List[int], + use_ue8m0: bool = False, use_bf16: bool = False, gran_k = 128): + assert get_mk_alignment_for_contiguous_layout() % gran_k == 0 + k = sum(ks) + + a = torch.randn((k, m), device='cuda', dtype=torch.bfloat16) + b = torch.randn((k, n), device='cuda', dtype=torch.bfloat16) + c = torch.randn((num_groups, m, n), device='cuda', dtype=torch.float) * 32 + d = c + ref_d = torch.empty_like(c) + + start = 0 + for i, group_k in enumerate(ks): + end = start + group_k + ref_d[i] = c[i] + (a[start:end].T @ b[start:end]) + start = end + + if use_bf16: + assert (major_a, major_b) == (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor) + return k, a, b, c, d, ref_d + + a_fp8 = per_channel_cast_to_fp8(a, use_ue8m0=use_ue8m0, gran_k=gran_k) + b_fp8 = per_channel_cast_to_fp8(b, use_ue8m0=use_ue8m0, gran_k=gran_k) + + # Transpose for K Major A/B + if (major_a, major_b) == (MajorTypeAB.KMajor, MajorTypeAB.KMajor): + a, sfa = a_fp8 + b, sfb = b_fp8 + new_a = torch.empty((sum(ks) * m, ), dtype=a.dtype, device=a.device) + new_b = torch.empty((sum(ks) * n, ), dtype=b.dtype, device=b.device) + prefix = 0 + for K in ks: + new_a[prefix * m : (prefix + K) * m] = a[prefix : prefix + K, ].T.flatten() + new_b[prefix * n : (prefix + K) * n] = b[prefix : prefix + K, ].T.flatten() + prefix += K + a_fp8, b_fp8 = (new_a, sfa.T), (new_b, sfb.T) + else: + assert (major_a, major_b) == (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor) + + return k, a_fp8, b_fp8, c, d, ref_d diff --git a/sgl_deep_gemm/tests/test_attention.py b/sgl_deep_gemm/tests/test_attention.py new file mode 100644 index 0000000000..479da5b56f --- /dev/null +++ b/sgl_deep_gemm/tests/test_attention.py @@ -0,0 +1,397 @@ +import dataclasses +import random +import torch +from typing import Tuple, List + +import deep_gemm +from deep_gemm.testing import ( + bench_kineto, + calc_diff, count_bytes, + ignore_env, get_arch_major, + test_filter +) +from deep_gemm.utils import ceil_div, per_custom_dims_cast_to_fp8, per_token_cast_to_fp4, cast_back_from_fp4 + +from generators import get_arch_major, generate_normal, get_ue8m0_usage, get_kernel_types, reset_seed, MajorTypeAB + + +def apply_skip_head_mid(d: torch.Tensor, head_splits: Tuple[int, int, int]): + left, mid, right = head_splits + m, n = d.shape + assert n % (left + right) == 0 + num_heads = n // (left + right) + + # Split and insert padding tensor + d = d.view(m, num_heads, -1) + d_left = d[:, :, :left] + d_right = d[:, :, -right:] + + d_mid = torch.zeros((m, num_heads, mid), dtype=d.dtype, device=d.device) + return torch.cat([d_left, d_mid, d_right], dim=2).view(m, -1) + + +def test_gemm_skip_head_mid() -> None: + print('Testing GEMM skip head mid:') + head_splits = (128, 64, 128) + + major_a, major_b = MajorTypeAB.KMajor, MajorTypeAB.KMajor + out_dtype, accumulate = torch.bfloat16, False + + for kernel_type in get_kernel_types(dtype=torch.float8_e4m3fn): + for m in (128, 4096): + for n, k in [(32768, 512), (8192, 512)]: + kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D' + use_ue8m0 = get_ue8m0_usage(kernel_type) + disable_ue8m0_cast = not use_ue8m0 + + a, b, _, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0) + d = apply_skip_head_mid(d, head_splits) + ref_d = apply_skip_head_mid(ref_d, head_splits) + + deep_gemm.fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, disable_ue8m0_cast=disable_ue8m0_cast) + diff = calc_diff(d, ref_d) + assert diff < 0.001, f'{m=}, {n=}, {k=}, {kernel_opt}, {diff:.5f}' + + t = bench_kineto(lambda: deep_gemm.fp8_gemm_nt_skip_head_mid(a, b, d, head_splits, disable_ue8m0_cast=disable_ue8m0_cast), + 'gemm_', suppress_kineto_output=True) + print(f' > Perf (m={m:5}, n={n:5}, k={k:5}, {kernel_opt}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{(count_bytes(a, b, d)) / 1e9 / t:4.0f} GB/s') + print() + + +def ref_fp8_mqa_logits(q: torch.Tensor, kv: torch.Tensor, weights: torch.Tensor, + cu_seqlen_ks: torch.Tensor, cu_seqlen_ke: torch.Tensor, cost_only: bool = False): + seq_len_kv = kv.shape[0] + + if cost_only: + start = cu_seqlen_ks.clamp(min=0, max=seq_len_kv) + end = cu_seqlen_ke.clamp(min=0, max=seq_len_kv) + count_ones_per_row = (end - start).clamp(min=0) + return count_ones_per_row.sum() + + k = kv + q = q.float() + k = k.float() + + mask_lo = torch.arange(0, seq_len_kv, device='cuda')[None, :] >= cu_seqlen_ks[:, None] + mask_hi = torch.arange(0, seq_len_kv, device='cuda')[None, :] < cu_seqlen_ke[:, None] + mask = mask_lo & mask_hi + + score = torch.einsum('mhd,nd->hmn', q, k) + logits = (score.relu() * weights.unsqueeze(-1).transpose(0, 1)).sum(dim=0) + logits = logits.masked_fill(~mask, float('-inf')) + + cost = mask.sum() + return logits, cost + + +def test_mqa_logits(): + + # Helper functions + def generate_ks_ke_tests(seq_len: int, seq_len_kv: int, disable_cp: bool): + if disable_cp: + ks = torch.zeros(seq_len, dtype=torch.int, device='cuda') + ke = torch.arange(seq_len, dtype=torch.int, device='cuda') + (seq_len_kv - seq_len) + return ks, ke + assert seq_len_kv % seq_len == 0 and seq_len % 2 == 0 + chunk_size = seq_len // 2 + cp_size = seq_len_kv // seq_len + # Select an arbitrary CP rank + cp_id = cp_size // 3 + ks = torch.zeros(seq_len, dtype=torch.int, device='cuda') + ke = torch.zeros(seq_len, dtype=torch.int, device='cuda') + for i in range(chunk_size): + ke[i] = cp_id * chunk_size + i + ke[i + chunk_size] = (cp_size * 2 - 1 - cp_id) * chunk_size + i + return ks, ke + + def enumerate_mqa_logits(): + for is_fp4 in ((True, False) if get_arch_major() == 10 else (False, )): + for logits_dtype in (torch.float, torch.bfloat16): + for compressed_logits, clean_logits in [(False, True), (True, False)]: + for seq_len in (2048, 4096): + for seq_len_kv in (4096, 8192): + for num_heads, head_dim in [(64, 128)]: + for disable_cp in (False, True): + yield is_fp4, logits_dtype, compressed_logits, clean_logits, seq_len, seq_len_kv, num_heads, head_dim, disable_cp + + print('Testing FP8 MQA Logits:') + for is_fp4, logits_dtype, compressed_logits, clean_logits, seq_len, seq_len_kv, num_heads, head_dim, disable_cp in enumerate_mqa_logits(): + # Generate random inputs + q = torch.randn(seq_len, num_heads, head_dim, device='cuda', dtype=torch.bfloat16) + kv = torch.randn(seq_len_kv, head_dim, device='cuda', dtype=torch.bfloat16) + weights = torch.randn(seq_len, num_heads, device='cuda', dtype=torch.float32) + ks, ke = generate_ks_ke_tests(seq_len, seq_len_kv, disable_cp) + + # Calculate reference logits + ref_logits, ref_cost = ref_fp8_mqa_logits(q, kv, weights, ks, ke) + + # Quantize Q and KV to FP4 / FP8 + if is_fp4: + q_fp4 = per_token_cast_to_fp4(q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + q_in = (q_fp4[0].view(seq_len, num_heads, head_dim // 2), q_fp4[1].view(seq_len, num_heads)) + q_simulated = cast_back_from_fp4(q_fp4[0], q_fp4[1], gran_k=32, use_packed_ue8m0=True).view(seq_len, num_heads, head_dim).to(torch.bfloat16) + + kv_fp4 = per_token_cast_to_fp4(kv.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + kv_in = (kv_fp4[0].view(seq_len_kv, head_dim // 2), kv_fp4[1].view(seq_len_kv)) + kv_simulated = cast_back_from_fp4(kv_fp4[0], kv_fp4[1], gran_k=32, use_packed_ue8m0=True).view(seq_len_kv, head_dim).to(torch.bfloat16) + else: + q_in = q.to(torch.float8_e4m3fn), None + q_simulated = q_in[0].to(torch.bfloat16) + kv_in = per_custom_dims_cast_to_fp8(kv, (0, ), False) + kv_simulated = (kv_in[0].float() * kv_in[1].unsqueeze(1)).to(torch.bfloat16) + + # Calculate reference logits + simulated_logits, _ = ref_fp8_mqa_logits(q_simulated, kv_simulated, weights, ks, ke) + + # Prepare kwargs + kernel_kwargs = dict( + q=q_in, kv=kv_in, weights=weights, + cu_seq_len_k_start=ks, cu_seq_len_k_end=ke, + clean_logits=clean_logits, max_seqlen_k=0, + logits_dtype=logits_dtype + ) + if compressed_logits: + max_seqlen_k = (ke - ks).max().item() + kernel_kwargs['max_seqlen_k'] = max_seqlen_k + + # Run kernel + logits = deep_gemm.fp8_fp4_mqa_logits(**kernel_kwargs) + + # Post process for compressed logits + if compressed_logits: + assert logits.size() == (seq_len, max_seqlen_k) + tmp = torch.full((seq_len, seq_len_kv), float('-inf'), device='cuda') + for i in range(seq_len): + tmp[i, ks[i] : ke[i]] = logits[i, : ke[i] - ks[i]] + logits = tmp + + # Validation + ref_neginf_mask = (ref_logits == float('-inf')) + neginf_mask = (logits == float('-inf')) + assert torch.equal(neginf_mask, ref_neginf_mask) + + ref_logits = ref_logits.masked_fill(ref_neginf_mask, 0) + simulated_logits = simulated_logits.masked_fill(ref_neginf_mask, 0) + logits = logits.masked_fill(ref_neginf_mask, 0) + diff = calc_diff(logits, ref_logits) + simulated_diff = calc_diff(logits, simulated_logits) + assert diff < 0.02 if is_fp4 else 1e-3, f"Diff: {diff}" + assert simulated_diff < 5e-6, f"Simulated Diff: {simulated_diff}" + + # Profiling + tflops = 2 * ref_cost * num_heads * head_dim / 1e12 + t, clean_t = bench_kineto(lambda: deep_gemm.fp8_fp4_mqa_logits(**kernel_kwargs), ('mqa_logits', 'clean_logits')) + clean_bytes = (seq_len * seq_len_kv - ref_cost) * 4 + count_bytes(ks, ke) + + print(f' > FP4={is_fp4}, BF16={logits_dtype == torch.bfloat16}, S={seq_len:4}, SKV={seq_len_kv:6}, H={num_heads:3}, D={head_dim:3}, CP={0 if disable_cp else 1}: ' + f'{tflops / t:4.0f} TFLOPS, {t * 1e6:4.0f} us, ' + f'{(count_bytes(q_in, kv_in, weights, ks, ke) + ref_cost * 4) / t / 1e9:4.0f} GB/s', end='') + print(f' | clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s' if clean_logits else '') + print() + + +def ref_paged_mqa_logits(q: torch.Tensor, kv_cache: torch.Tensor, + weights: torch.Tensor, context_lens: torch.Tensor, block_tables: torch.Tensor, + max_model_len: int, use_2d_context_lens: bool): + batch_size, next_n, num_heads, dim = q.size() + num_block, block_size, _, dim = kv_cache.size() + logits = torch.full([batch_size * next_n, max_model_len], float('-inf'), device=q.device, dtype=torch.float32) + context_lens = context_lens.tolist() + for i in range(batch_size): + context_len = context_lens[i] + q_offsets = torch.full((next_n, ), context_len, device='cuda', dtype=torch.int32) if use_2d_context_lens \ + else torch.arange(context_len - next_n, context_len, device='cuda') + weight_slice = weights[i * next_n:(i + 1) * next_n, :].transpose(0, 1).contiguous() + + num_blocks = (context_len + block_size - 1) // block_size + block_idxs = block_tables[i][:num_blocks] + kv_slice = kv_cache[block_idxs] # [num_blocks, block_size, kv_heads, dim] + kx = kv_slice.permute(2, 3, 0, 1).reshape(kv_slice.size(2), dim, -1) # [kv_heads, dim, total_tokens] + qx = q[i].transpose(0, 1) # q[i]: [next_n, num_heads, dim] -> [num_heads, next_n, dim] + s = torch.matmul(qx, kx).to(logits.dtype) # [num_heads, next_n, dim] @ [1, dim, total_tokens] -> [num_heads, next_n, total_tokens] + + total_len = num_blocks * block_size + k_offsets = torch.arange(0, total_len, device=q.device) + mask = (k_offsets[None, :] < context_len) & (k_offsets[None, :] <= q_offsets[:, None]) + s = torch.where(mask[None, :, :], s, float('-inf')) # mask shape: [1, next_n, total_tokens] + s = torch.relu(s) * weight_slice[..., None] # weight_slice: [num_heads, next_n] -> [num_heads, next_n, 1] + s = s.sum(dim=0) # [next_n, total_tokens] + logits[i * next_n:(i + 1) * next_n, :total_len] = torch.where(k_offsets[None, :] <= q_offsets[:, None], s, float('-inf')) + + return logits + + +def test_paged_mqa_logits(): + + # Helper functions + def kv_cache_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + num_blocks, block_size, num_heads, head_dim = x.shape + assert num_heads == 1 + x_amax = x.abs().float().amax(dim=3, keepdim=True).clamp(1e-4) + sf = x_amax / 448.0 + x_scaled = (x * (1.0 / sf)).to(torch.float8_e4m3fn) + x_cast_back = x_scaled.float() * sf + + x_fp8 = torch.empty((num_blocks, block_size * (head_dim + 4)), device=x.device, dtype=torch.uint8) + x_fp8[ :, : block_size * head_dim] = x_scaled.view(num_blocks, block_size * head_dim).view(torch.uint8) + x_fp8[ :, block_size * head_dim :] = sf.view(num_blocks, block_size).view(torch.uint8) + return x_fp8.view(num_blocks, block_size, num_heads, head_dim + 4), x_cast_back.to(x.dtype) + + def kv_cache_cast_to_fp4(x: torch.Tensor) -> torch.Tensor: + num_blocks, block_size, num_heads, head_dim = x.shape + assert num_heads == 1 and head_dim == 128 + x_scaled, sf = per_token_cast_to_fp4(x.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + x_cast_back = cast_back_from_fp4(x_scaled, sf, gran_k=32, use_packed_ue8m0=True).view(num_blocks, block_size, 1, head_dim) + + x_fp4 = torch.empty((num_blocks, block_size * (head_dim // 2 + 4)), device=x.device, dtype=torch.uint8) + x_fp4[ :, : block_size * head_dim // 2] = x_scaled.view(num_blocks, block_size * head_dim // 2).view(torch.uint8) + x_fp4[ :, block_size * head_dim // 2 :] = sf.view(num_blocks, block_size).view(torch.uint8) + return x_fp4.view(num_blocks, block_size, num_heads, head_dim // 2 + 4), x_cast_back.to(x.dtype) + + def enumerate_paged_mqa_logits(): + arch_major = get_arch_major() + for is_varlen in ((True, False) if arch_major == 10 else (False, )): + for is_fp4 in ((True, False) if arch_major == 10 else (False, )): + for logits_dtype in (torch.float, torch.bfloat16): + for block_kv in ((32, 64) if arch_major == 10 else (64, )): + for use_2d_context_lens, clean_logits in [(True, False)]: + for batch_size in (256, ): + for next_n in ((1, ) if is_varlen else ((1, 2, 4, 5, 6) if arch_major == 10 else (1, 2))): + for max_tokens_per_batch in ((1, 4, 10) if is_varlen else (1, )): + for num_heads, head_dim in [(64, 128)]: + for avg_kv in (8192, 32768): + yield is_varlen, is_fp4, logits_dtype, block_kv, use_2d_context_lens, clean_logits, batch_size, next_n, max_tokens_per_batch, num_heads, head_dim, avg_kv + + + print('Testing FP8/FP4 Paged MQA Logits:') + max_model_len = 111 * 1024 + num_total_blocks = max_model_len * 5 + + for is_varlen, is_fp4, logits_dtype, block_kv, use_2d_context_lens, clean_logits, batch_size, next_n, max_tokens_per_batch, num_heads, head_dim, avg_kv in enumerate_paged_mqa_logits(): + # Varlen: flatten raw_batch_size sequences with variable tokens into (batch_size, 1, ...) + raw_batch_size, raw_next_n = batch_size, next_n + if is_varlen: + tokens_per_seq = torch.randint(1, max_tokens_per_batch + 1, (raw_batch_size,), device='cuda', dtype=torch.int) + indices = torch.arange(raw_batch_size, device='cuda', dtype=torch.int).repeat_interleave(tokens_per_seq) + batch_size, next_n = tokens_per_seq.sum().item(), 1 + else: + tokens_per_seq, indices = None, None + + # Generate random inputs + q = torch.randn((batch_size, next_n, num_heads, head_dim), device='cuda', dtype=torch.bfloat16) + kv_cache = torch.randn((num_total_blocks, block_kv, 1, head_dim), device='cuda', dtype=torch.bfloat16) + weights = torch.randn((batch_size * next_n, num_heads), device='cuda', dtype=torch.float) + context_lens = torch.randint(int(0.7 * avg_kv), int(1.3 * avg_kv), (raw_batch_size,), device='cuda', dtype=torch.int) + + if is_varlen: + max_ctx_len_per_seq = context_lens + (tokens_per_seq - 1) + else: + max_ctx_len_per_seq = context_lens + + # Assign block tables (per-sequence, sized by the largest ctx_len within the sequence) + seq_sum_lens = context_lens.sum().item() + num_blocks_per_query = ceil_div(max_ctx_len_per_seq, block_kv) + block_table = torch.empty((raw_batch_size, num_blocks_per_query.max().item()), device='cuda', dtype=torch.int) + block_idx_pool = torch.randperm(num_total_blocks, device='cuda', dtype=torch.int) + offset = 0 + for i, num_blocks in enumerate(num_blocks_per_query.tolist()): + block_table[i, :num_blocks] = block_idx_pool[offset : offset + num_blocks] + offset += num_blocks + if is_varlen: + context_lens = context_lens.repeat_interleave(tokens_per_seq) + offsets_within_seq = torch.cat([ + torch.arange(n.item(), device='cuda', dtype=torch.int) + for n in tokens_per_seq + ]) + context_lens = context_lens + offsets_within_seq + block_table = block_table.repeat_interleave(tokens_per_seq, dim=0) + + # Calculate reference logits + ref_logits = ref_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, max_model_len, use_2d_context_lens) + + # Quantize Q and KV cache to FP4 / FP8 + if is_fp4: + q_fp4 = per_token_cast_to_fp4(q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + q_in = (q_fp4[0].view(batch_size, next_n, num_heads, head_dim // 2), q_fp4[1].view(batch_size, next_n, num_heads)) + q_simulated = cast_back_from_fp4(q_fp4[0], q_fp4[1], gran_k=32, use_packed_ue8m0=True).view(batch_size, next_n, num_heads, head_dim).to(torch.bfloat16) + kv_in, kv_simulated = kv_cache_cast_to_fp4(kv_cache) + else: + q_in = q.to(torch.float8_e4m3fn), None + q_simulated = q_in[0].to(torch.bfloat16) + kv_in, kv_simulated = kv_cache_cast_to_fp8(kv_cache) + + # Calculate simulated reference logits + simulated_logits = ref_paged_mqa_logits(q_simulated, kv_simulated, weights, context_lens, block_table, max_model_len, use_2d_context_lens) + + # Prepare masks and context lengths with NextN + positions = torch.arange(max_model_len, device='cuda').unsqueeze(0).expand(batch_size * next_n, -1) + if use_2d_context_lens: + if is_varlen: + # Varlen: context_lens is already per-token (shape [total_tokens]); + # just reshape to (total_tokens, 1) so each token keeps its own ctx_len. + context_lens_nextn = context_lens.view(-1, 1) + else: + context_lens_nextn = ((context_lens.unsqueeze(1) + 1) * torch.rand(batch_size, next_n, device='cuda')).int() + # Ensure last token matches actual length + context_lens_nextn[:, -1] = context_lens + ref_neginf_mask = ~(positions < context_lens_nextn.view(-1, 1)) + else: + context_lens_nextn = context_lens + offsets = torch.arange(batch_size * next_n, device='cuda') + limits = (context_lens[offsets // next_n] - next_n + offsets % next_n).unsqueeze(1) + ref_neginf_mask = ~(positions <= limits) + + # Run Kernel + kernel_kwargs = dict( + q=q_in, kv_cache=kv_in, weights=weights, + context_lens=context_lens_nextn, block_table=block_table, + schedule_meta=deep_gemm.get_paged_mqa_logits_metadata(context_lens_nextn, block_kv, deep_gemm.get_num_sms(), indices=indices), + max_context_len=max_model_len, clean_logits=clean_logits, logits_dtype=logits_dtype, + indices=indices, + ) + logits = deep_gemm.fp8_fp4_paged_mqa_logits(**kernel_kwargs) + + # Validation + assert logits.dtype == logits_dtype + logits = logits.to(torch.float) + + if clean_logits: + assert torch.equal(logits == float('-inf'), ref_neginf_mask), "Mask mismatch" + + logits_masked = logits.masked_fill(ref_neginf_mask, 0) + ref_masked = ref_logits.masked_fill(ref_neginf_mask, 0) + simulated_masked = simulated_logits.masked_fill(ref_neginf_mask, 0) + diff = calc_diff(logits_masked, ref_masked) + simulated_diff = calc_diff(logits_masked, simulated_masked) + assert diff < 0.02 if is_fp4 else 1e-3, f"Diff: {diff}" + assert simulated_diff < 5e-6, f"Simulated Diff: {simulated_diff}" + + # Profiling + sum_lens = context_lens.sum().item() + tflops_calc = 2 * sum_lens * next_n * num_heads * head_dim / 1e12 + kv_bytes_per_token = head_dim / (2 if is_fp4 else 1) + 4 + # KV is read once per sequence; for varlen sum_lens overcounts (per-token), so use seq_sum_lens + kv_sum_lens = seq_sum_lens if is_varlen else sum_lens + total_bytes = count_bytes(q, weights) + kv_sum_lens * kv_bytes_per_token + (sum_lens * next_n * logits_dtype.itemsize) + + t, clean_t = bench_kineto(lambda: deep_gemm.fp8_fp4_paged_mqa_logits(**kernel_kwargs), ('paged_mqa_logits', 'clean_logits')) + print(f' > FP4={is_fp4}, BF16={logits_dtype == torch.bfloat16}, BLOCK_KV={block_kv}, BSZ={raw_batch_size:3}, NextN={raw_next_n:1}, H={num_heads:2}, D={head_dim:2}, L={avg_kv:6}: ' + f'{tflops_calc / t:4.0f} TFLOPS, {t * 1e6:3.0f} us, {total_bytes / t / 1e9:4.0f} GB/s', end='') + if is_varlen: + print(f' | Varlen, MaxTPB={max_tokens_per_batch}, NumTokens={batch_size}', end='') + print(f' | clean: {clean_t*1e6:3.0f} us' if clean_logits else '') + print() + + + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + + test_gemm_skip_head_mid() + test_mqa_logits() + test_paged_mqa_logits() diff --git a/sgl_deep_gemm/tests/test_bf16.py b/sgl_deep_gemm/tests/test_bf16.py new file mode 100644 index 0000000000..703337b86e --- /dev/null +++ b/sgl_deep_gemm/tests/test_bf16.py @@ -0,0 +1,223 @@ +import copy +import numpy as np +import random +import torch + +import deep_gemm +from deep_gemm.testing import ( + bench_kineto, + calc_diff, count_bytes +) +from generators import ( + get_arch_major, layout_masked_to_psum, align, + enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous, + generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous, + get_mk_alignment_for_contiguous_layout +) + + +def test_gemm() -> None: + print('Testing GEMM:') + scores = [] + for kernel_type, _, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(torch.bfloat16): + major_opt = 'N' if major_a.is_k_major() else 'T' + major_opt += 'T' if major_b.is_k_major() else 'N' + out_opt = 'FP32' if out_dtype == torch.float else 'BF16' + acc_opt = f'acc={int(accumulate)}' + + for test_alias in (False, True): + a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_bf16=True) + func_name = f'bf16_gemm_{major_opt.lower() if test_alias else "nt"}' + if test_alias: + a = a if major_a.is_k_major() else a.T + b = b if major_b.is_k_major() else b.T + assert a.is_contiguous() and b.is_contiguous() + getattr(deep_gemm, func_name)(a, b, d, c=c) + diff = calc_diff(d, ref_d) + assert diff < 1e-5, (f'{m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=}, ' + f'{diff:.5f}, alias={test_alias}') + a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_bf16=True) + + t = bench_kineto(lambda: deep_gemm.bf16_gemm_nt(a, b, d, c=c), 'bf16_gemm', suppress_kineto_output=True) + cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c), ('nvjet', 'reduce'), suppress_kineto_output=True) + print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, layout={major_opt}, {out_opt}, {acc_opt}): ' + f'{t * 1e6:7.1f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{(count_bytes(a, b, d) + count_bytes(c) * int(accumulate)) / 1e9 / t:4.0f} GB/s | ' + f'{(cublas_t + split_k_t) / t:.2f}x cuBLAS') + if cublas_t > 0: + scores.append((cublas_t + split_k_t) / t) + print(f"Average speedup over cuBLASLt: {float(np.prod(scores)) ** (1.0 / len(scores)):.3f}x\n") + + +def test_m_grouped_gemm_contiguous() -> None: + print('Testing m-grouped contiguous GEMM:') + + for _, _, num_groups, expected_m_per_group, n, k, major_a, major_b, use_psum_layout in enumerate_m_grouped_contiguous(torch.bfloat16): + major_opt = 'N' if major_a.is_k_major() else 'T' + major_opt += 'T' if major_b.is_k_major() else 'N' + + # Select best alignment + alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout() + deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) + + for test_alias in (False, True): + m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, + use_bf16=True, use_psum_layout=use_psum_layout) + func_name = f"m_grouped_bf16_gemm_{(major_opt.lower() if test_alias else 'nt')}_contiguous" + if test_alias: + assert major_a.is_k_major() + b = b if major_b.is_k_major() else b.mT + assert a[0].is_contiguous() and b[0].is_contiguous() + getattr(deep_gemm, func_name)(a, b, d, grouped_layout, use_psum_layout=use_psum_layout) + if use_psum_layout: + for j in range(num_groups): + start = 0 if j == 0 else align(grouped_layout[j - 1], get_mk_alignment_for_contiguous_layout()) + end = grouped_layout[j] + diff = calc_diff(d[start : end], ref_d[start : end]) + assert diff < 1e-5, f'{m=}, {n=}, {k=}, {major_opt}, {diff:.5f}, alias={test_alias}' + else: + diff = calc_diff(d, ref_d) + assert diff < 1e-5, f'{m=}, {n=}, {k=}, {major_opt}, {diff:.5f}, alias={test_alias}' + m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, + use_bf16=True, use_psum_layout=use_psum_layout) + + # noinspection PyShadowingNames + def test_func(): + deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, use_psum_layout=use_psum_layout) + + t = bench_kineto(test_func, 'bf16_gemm', suppress_kineto_output=True) + print(f' > Perf ({num_groups=}, m={m:5}, n={n:5}, k={k:5}, layout={major_opt}, psum={use_psum_layout}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s') + print() + + +def test_m_grouped_gemm_masked() -> None: + print('Testing m-grouped masked GEMM:') + + # TODO: when the actual `m` is greater than `expected_m_per_group`, efficiency may significantly decrease. + for _, _, num_groups, max_m, expected_m_per_group, n, k, use_psum_layout in enumerate_m_grouped_masked(torch.bfloat16): + num_tests = 8 + sum_t, max_t = 0, 0 + sum_ops, sum_bytes = 0, 0 + + # Select best alignment + alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(int(expected_m_per_group * 1.2)) + deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) + + for i in range(num_tests): + a, b, masked_m, psum_m, d, ref_d = generate_m_grouped_masked(num_groups, max_m, expected_m_per_group, n, k, + use_bf16=True, use_psum_layout=use_psum_layout) + if use_psum_layout: + a_psum = layout_masked_to_psum(a, psum_m) + d_psum = layout_masked_to_psum(d, psum_m) + + # noinspection PyShadowingNames + def test_func(): + if use_psum_layout: + deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a_psum, b, d_psum, psum_m, + use_psum_layout=True, expected_m_for_psum_layout=expected_m_per_group) + else: + deep_gemm.m_grouped_bf16_gemm_nt_masked(a, b, d, masked_m, expected_m_per_group) + + test_func() + for j in range(num_groups): + if masked_m[j].item() == 0: + continue + if use_psum_layout: + d_slice = d_psum[: psum_m[j]] if j == 0 else d_psum[align(psum_m[j - 1], get_mk_alignment_for_contiguous_layout()): psum_m[j]] + else: + d_slice = d[j, :masked_m[j].item()] + diff = calc_diff(d_slice, ref_d[j, :masked_m[j].item()]) + assert diff < 1e-5, f'{max_m=}, {n=}, {k=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}' + + + # Test performance with fixed shapes + valid_m = masked_m.sum().item() + t = bench_kineto(test_func, 'bf16_gemm', suppress_kineto_output=True) + + sum_t += t + max_t = max(max_t, t) + sum_ops += 2 * valid_m * n * k + sum_bytes += count_bytes(a, d) * valid_m / (max_m * num_groups) + count_bytes(b) + + print(f' > Perf (num_groups={num_groups:2}, expected_m_per_group={expected_m_per_group:4}, n={n:4}, k={k:4}, ' + f'psum={1 if use_psum_layout else 0}): ' + f'{sum_t / num_tests * 1e6:4.0f} us (max: {max_t * 1e6:3.0f} us) | ' + f'{sum_ops / sum_t / 1e12:4.0f} TFLOPS | ' + f'{sum_bytes / sum_t / 1e9:4.0f} GB/s') + print() + + +def test_k_grouped_gemm_contiguous() -> None: + print('Testing k-grouped contiguous GEMM:') + + # TODO: Support arbitrary alignment + deep_gemm.set_mk_alignment_for_contiguous_layout(128) + + for num_groups, m, n, major_a, major_b, ks, expected_k_per_group in enumerate_k_grouped_contiguous(torch.bfloat16): + for test_empty_groups in (False, True): + new_ks = copy.deepcopy(ks) + if test_empty_groups and len(ks) > 1: + new_ks[random.randint(0, num_groups - 1)] = 0 + k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, new_ks, use_bf16=True) + new_ks_tensor = torch.tensor(new_ks, dtype=torch.int, device='cuda') + deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, new_ks, new_ks_tensor, c) + + diff = calc_diff(d, ref_d) + assert diff < 1e-5, f'{m=}, {n=}, {k=}, {ks=}, {diff:.7f}' + + # Test performance + k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_bf16=True) + ks_tensor = torch.tensor(ks, dtype=torch.int, device='cuda') + + # noinspection PyShadowingNames + def test_func(): + deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, ks_tensor, c) + + t = bench_kineto(test_func, 'bf16_gemm', suppress_kineto_output=True) + print(f' > Perf ({num_groups=:2}, m={m:5}, n={n:5}, k={k:5}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, c, d) / 1e9 / t:4.0f} GB/s') + print() + + +def test_cublaslt_gemm() -> None: + print('Testing cuBLASLt GEMM:') + for kernel_type, _, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(dtype=torch.bfloat16): + major_opt = 'N' if major_a.is_k_major() else 'T' + major_opt += 'T' if major_b.is_k_major() else 'N' + out_opt = 'FP32' if out_dtype == torch.float else 'BF16' + acc_opt = f'acc={int(accumulate)}' + + a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_bf16=True) + deep_gemm.cublaslt_gemm_nt(a, b, d, c) + diff = calc_diff(d, ref_d) + assert diff < 6e-7, f'{diff=}, ({m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=})' + + t_nvjet, t_gemv, t_gemm = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c), ('nvjet', 'gemv', 'gemm'), suppress_kineto_output=True) + t = t_nvjet + t_gemv + t_gemm + print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, layout={major_opt}, {out_opt}, {acc_opt}): ' + f'{t * 1e6:5.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{(count_bytes(a, b, d) + count_bytes(c) * int(accumulate)) / 1e9 / t:4.0f} GB/s') + print() + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + + print('Library path:') + print(f' > {deep_gemm.__path__}\n') + + if get_arch_major() >= 9: + test_gemm() + test_m_grouped_gemm_contiguous() + test_m_grouped_gemm_masked() + test_k_grouped_gemm_contiguous() + + test_cublaslt_gemm() diff --git a/sgl_deep_gemm/tests/test_einsum.py b/sgl_deep_gemm/tests/test_einsum.py new file mode 100644 index 0000000000..57f5459288 --- /dev/null +++ b/sgl_deep_gemm/tests/test_einsum.py @@ -0,0 +1,181 @@ +import random +import torch + +import deep_gemm +from deep_gemm.testing import ( + bench_kineto, + calc_diff, count_bytes, + get_arch_major, test_filter +) +from deep_gemm.utils.math import ( + ceil_div, + per_block_cast_to_fp8, per_channel_cast_to_fp8, per_token_cast_to_fp8 +) + + +def test_bmk_bnk_mn() -> None: + print('Testing "bmk, bnk -> mn":') + for s in (129, 4096, 8192): + for m, n, k in [(128, 384, 128), (256, 256, 256), (384, 128, 384)]: + for dtype in (torch.float, torch.bfloat16): + a = torch.randn((s, m, k), dtype=torch.bfloat16, device='cuda') + b = torch.randn((s, n, k), dtype=torch.bfloat16, device='cuda') + d = torch.randn((m, n), dtype=dtype, device='cuda') + c = d if dtype == torch.float else None + + # Test correctness + ref_d = (c if dtype == torch.float else 0) + torch.bmm(a.float(), b.float().mT).sum(0) + deep_gemm.einsum('bmk,bnk->mn', a, b, d, c=c) + assert calc_diff(d, ref_d) < 1e-5 + + t = bench_kineto(lambda: deep_gemm.einsum('bmk,bnk->mn', a, b, d, c=c), 'bmn_bnk_mn_gemm_impl', suppress_kineto_output=True) + print(f' > Perf (b={s:4.0f}, {m=}, {n=}, {k=}, {"FP32" if dtype == torch.float else "BF16"}): ', + f'{t * 1e6:4.0f} us | ' + f'{2 * s * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{(count_bytes(a, b) + (d.numel() * 4)) / 1e9 / t:4.0f} GB/s') + print() + + +def test_bhr_hdr_bhd(): + print('Testing "bhr, hdr -> bhd":') + for h, r, d in [(128, 512, 128), (8, 4096, 1024)]: + for b in (4, 32, 128, 4096, 8192): + x = torch.randn((b, h, r), device='cuda', dtype=torch.bfloat16) + fy = torch.randn((h, d, r + 128), device='cuda', dtype=torch.bfloat16) + y = fy[:, :, :r] + ref_z = torch.einsum('bhr,hdr->bhd', x, y) + z = torch.empty((b, h, d), device='cuda', dtype=torch.bfloat16) + deep_gemm.einsum('bhr,hdr->bhd', x, y, z) + assert calc_diff(z, ref_z) < 1e-10 + + t = bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z), 'gemm', suppress_kineto_output=True) + t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z, use_cublaslt=True), 'nvjet', suppress_kineto_output=True) + print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ', + f'{t * 1e6:4.0f} us | ' + f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes((x, y, z)) / t / 1e9:4.0f} GB/s | ' + f'{t_cublaslt / t:4.2f} x') + print() + + +def test_bhd_hdr_bhr(): + print('Testing "bhd, hdr -> bhr":') + for h, r, d in [(128, 512, 128), (8, 4096, 1024)]: + for b in (4, 32, 128, 4096, 8192): + x = torch.randn((b, h, d), device='cuda', dtype=torch.bfloat16) + fy = torch.randn((h, d, r + 128), device='cuda', dtype=torch.bfloat16) + y = fy[:, :, :r] + ref_z = torch.einsum('bhd,hdr->bhr', x, y) + z = torch.empty((b, h, r), device='cuda', dtype=torch.bfloat16) + deep_gemm.einsum('bhd,hdr->bhr', x, y, z) + assert calc_diff(z, ref_z) < 1e-10 + + t = bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z), 'gemm', suppress_kineto_output=True) + t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z, use_cublaslt=True), 'nvjet', suppress_kineto_output=True) + print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ', + f'{t * 1e6:4.0f} us | ' + f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes((x, y, z)) / t / 1e9:4.0f} GB/s | ' + f'{t_cublaslt / t:4.2f} x') + print() + + +def test_fp8_bhr_hdr_bhd(use_ue8m0: bool = True): + print('Testing FP8 "bhr, hdr -> bhd":') + for h, r, d in [(8, 4096, 1024)]: + for b in (4, 32, 128, 4096, 8192): + x = torch.randn((b, h, r), device='cuda', dtype=torch.bfloat16) + y = torch.randn((h, d, r), device='cuda', dtype=torch.bfloat16) + ref_z = torch.einsum('bhr,hdr->bhd', x, y) + + x_fp8 = per_token_cast_to_fp8(x.view(-1, r), use_ue8m0=use_ue8m0) + x_fp8 = x_fp8[0].view(b, h, r), x_fp8[1].view(b, h, ceil_div(r, 128)) + y_fp8 = (torch.empty_like(y, dtype=torch.float8_e4m3fn), + torch.empty((h, ceil_div(d, 128), ceil_div(r, 128)), device='cuda', dtype=torch.float)) + for i in range(h): + y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i], use_ue8m0=use_ue8m0) + z = torch.empty((b, h, d), device='cuda', dtype=torch.bfloat16) + + deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, z) + assert calc_diff(z, ref_z) < 1e-3 + + t = bench_kineto(lambda: deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, z), 'gemm_', suppress_kineto_output=True) + t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z, use_cublaslt=True), 'nvjet', suppress_kineto_output=True) + print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ', + f'{t * 1e6:4.0f} us | ' + f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes((x_fp8, y_fp8, z)) / t / 1e9:4.0f} GB/s | ' + f'{t_cublaslt / t:4.2f} x') + print() + + +@test_filter(lambda: get_arch_major() >= 10) +def test_fp8_bhd_hdr_bhr(use_ue8m0: bool = True): + print('Testing FP8 "bhd, hdr -> bhr":') + for h, r, d in [(8, 4096, 1024)]: + for b in (4, 32, 128, 4096, 8192): + x = torch.randn((b, h, d), device='cuda', dtype=torch.bfloat16) + y = torch.randn((h, d, r), device='cuda', dtype=torch.bfloat16) + ref_z = torch.einsum('bhd,hdr->bhr', x, y) + + x_fp8 = per_token_cast_to_fp8(x.view(-1, d), use_ue8m0=use_ue8m0) + x_fp8 = x_fp8[0].view(b, h, d), x_fp8[1].view(b, h, ceil_div(d, 128)) + y_fp8 = (torch.empty_like(y, dtype=torch.float8_e4m3fn), + torch.empty((h, ceil_div(d, 128), ceil_div(r, 128)), device='cuda', dtype=torch.float)) + for i in range(h): + y_fp8[0][i], y_fp8[1][i] = per_block_cast_to_fp8(y[i], use_ue8m0=use_ue8m0) + z = torch.empty((b, h, r), device='cuda', dtype=torch.bfloat16) + + deep_gemm.fp8_einsum('bhd,hdr->bhr', x_fp8, y_fp8, z) + assert calc_diff(z, ref_z) < 1e-3 + + t = bench_kineto(lambda: deep_gemm.fp8_einsum('bhd,hdr->bhr', x_fp8, y_fp8, z), 'gemm_', suppress_kineto_output=True) + t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z, use_cublaslt=True), 'nvjet', suppress_kineto_output=True) + print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ', + f'{t * 1e6:4.0f} us | ' + f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes((x_fp8, y_fp8, z)) / t / 1e9:4.0f} GB/s | ' + f'{t_cublaslt / t:4.2f} x') + print() + + +@test_filter(lambda: get_arch_major() >= 10) +def test_fp8_bhd_bhr_hdr(use_ue8m0: bool = True): + print('Testing FP8 "bhd, bhr -> hdr":') + for h, r, d in [(8, 4096, 1024)]: + for b in (4096, 8192): + x = torch.randn((b, h, d), device='cuda', dtype=torch.bfloat16) + y = torch.randn((b, h, r), device='cuda', dtype=torch.bfloat16) + z_0 = torch.randn((h, d, r), device='cuda', dtype=torch.float32) * 10 + ref_z = z_0 + torch.einsum('bhd,bhr->hdr', x, y) + + x_fp8 = per_channel_cast_to_fp8(x.view(b, -1), use_ue8m0=use_ue8m0) + y_fp8 = per_channel_cast_to_fp8(y.view(b, -1), use_ue8m0=use_ue8m0) + x_fp8 = (x_fp8[0].view(b, h, d), x_fp8[1].view(ceil_div(b, 128), h, d)) + y_fp8 = (y_fp8[0].view(b, h, r), y_fp8[1].view(ceil_div(b, 128), h, r)) + z = z_0.clone() + deep_gemm.fp8_einsum('bhd,bhr->hdr', x_fp8, y_fp8, z, z, recipe=(1, 1, 128)) + assert calc_diff(z, ref_z) < 1e-3 + + t = bench_kineto(lambda: deep_gemm.fp8_einsum('bhd,bhr->hdr', x_fp8, y_fp8, z, z, recipe=(1, 1, 128)), 'gemm_', suppress_kineto_output=True) + print(f' > Perf ({b=:4.0f}, {h=}, {r=}, {d=}): ', + f'{t * 1e6:4.0f} us | ' + f'{2 * b * h * r * d / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes((x_fp8, y_fp8, z, z)) / t / 1e9:4.0f} GB/s') + print() + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + + print('Library path:') + print(f' > {deep_gemm.__path__}\n') + + test_bmk_bnk_mn() + test_bhr_hdr_bhd() + test_bhd_hdr_bhr() + + test_fp8_bhr_hdr_bhd() + test_fp8_bhd_hdr_bhr() + test_fp8_bhd_bhr_hdr() diff --git a/sgl_deep_gemm/tests/test_fp8_fp4.py b/sgl_deep_gemm/tests/test_fp8_fp4.py new file mode 100644 index 0000000000..42038e459d --- /dev/null +++ b/sgl_deep_gemm/tests/test_fp8_fp4.py @@ -0,0 +1,223 @@ +import copy +import numpy as np +import random +import torch + +import deep_gemm +from deep_gemm.testing import ( + bench_kineto, + calc_diff, count_bytes, + ignore_env, get_arch_major +) + +from generators import ( + KernelType, get_ue8m0_usage, layout_masked_to_psum, align, + enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous, + generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous, + get_mk_alignment_for_contiguous_layout +) + + +def test_gemm() -> None: + print('Testing GEMM:') + scores = [] + for kernel_type, quant_config, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(torch.float8_e4m3fn): + major_opt = 'N' if major_a.is_k_major() else 'T' + major_opt += 'T' if major_b.is_k_major() else 'N' + out_opt = 'FP32' if out_dtype == torch.float else 'BF16' + acc_opt = f'acc={int(accumulate)}' + kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D' + use_ue8m0 = get_ue8m0_usage(kernel_type) + disable_ue8m0_cast = not use_ue8m0 + recipe, recipe_a, recipe_b = quant_config.get_recipes(is_wgrad=(kernel_type.is_1d1d() and accumulate)) + + for test_alias in (False, True): + a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0, quant_config=quant_config) + func_name = f'fp8_fp4_gemm_{major_opt.lower() if test_alias else "nt"}' + if test_alias: + a = a if major_a.is_k_major() else (a[0].T, a[1].T) + b = b if major_b.is_k_major() else (b[0].T, b[1].T) + assert a[0].is_contiguous() and b[0].is_contiguous() + a, a_sf = a + b, b_sf = b + getattr(deep_gemm, func_name)(a, a_sf, b, b_sf, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + diff = calc_diff(d, ref_d) + assert diff < quant_config.max_diff(), (f'{m=}, {n=}, {k=}, {kernel_opt}, {major_opt=}, {accumulate=}, {out_dtype=}, ' + f'{diff:.5f}, alias={test_alias}') + a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0, quant_config=quant_config) + t = bench_kineto(lambda: deep_gemm.fp8_fp4_gemm_nt(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b), + 'gemm_', suppress_kineto_output=True) + cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a[0], b[0], d, c=c), ('nvjet', 'reduce'), suppress_kineto_output=True) \ + if not quant_config.is_fp4_a and not quant_config.is_fp4_b else (0, 0) + print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, {kernel_opt}, layout={major_opt}, {out_opt}, {acc_opt}): ' + f'{t * 1e6:6.1f} us | {2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{(count_bytes(a, b, d) + count_bytes(c) * int(accumulate)) / 1e9 / t:4.0f} GB/s | ' + f'{(cublas_t + split_k_t) / t:.2f}x cuBLAS') + if cublas_t > 0: + scores.append((cublas_t + split_k_t) / t) + print(f"Average FP8xFP8 GEMM speedup over cuBLASLt: {float(np.prod(scores)) ** (1.0 / len(scores)):.3f}x\n") + + +def test_m_grouped_gemm_contiguous() -> None: + print('Testing m-grouped contiguous GEMM:') + + for kernel_type, quant_config, num_groups, expected_m_per_group, n, k, major_a, major_b, use_psum_layout in enumerate_m_grouped_contiguous(dtype=torch.float8_e4m3fn): + major_opt = 'N' if major_a.is_k_major() else 'T' + major_opt += 'T' if major_b.is_k_major() else 'N' + kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D' + use_ue8m0 = get_ue8m0_usage(kernel_type) + disable_ue8m0_cast = not use_ue8m0 + recipe, recipe_a, recipe_b = quant_config.get_recipes() + + # Select best alignment + alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout() + deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) + + for test_alias in (False, True): + m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, + use_ue8m0=use_ue8m0, use_psum_layout=use_psum_layout, + quant_config=quant_config) + func_name = f"m_grouped_fp8_fp4_gemm_{(major_opt.lower() if test_alias else 'nt')}_contiguous" + if test_alias: + assert major_a.is_k_major() + b = b if major_b.is_k_major() else (b[0].mT, b[1].mT) + assert a[0].is_contiguous() and b[0].is_contiguous() + getattr(deep_gemm, func_name)(a, b, d, grouped_layout, disable_ue8m0_cast=disable_ue8m0_cast, use_psum_layout=use_psum_layout, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + if use_psum_layout: + for j in range(num_groups): + start = 0 if j == 0 else align(grouped_layout[j - 1], get_mk_alignment_for_contiguous_layout()) + end = grouped_layout[j] + diff = calc_diff(d[start : end], ref_d[start : end]) + assert diff < quant_config.max_diff(), f'{m=}, {n=}, {k=}, {major_opt}, {kernel_opt}, {diff:.5f}, alias={test_alias}' + else: + diff = calc_diff(d, ref_d) + assert diff < quant_config.max_diff(), f'{m=}, {n=}, {k=}, {major_opt}, {kernel_opt}, {diff:.5f}, alias={test_alias}' + m, a, b, grouped_layout, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, + use_ue8m0=use_ue8m0, use_psum_layout=use_psum_layout, + quant_config=quant_config) + + # noinspection PyShadowingNames + def test_func(): + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(a, b, d, grouped_layout, disable_ue8m0_cast=disable_ue8m0_cast, use_psum_layout=use_psum_layout, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + + t = bench_kineto(test_func, 'gemm_', suppress_kineto_output=True) + print(f' > Perf ({num_groups=}, m={m:5}, n={n:6}, k={k:5}, {kernel_opt}, layout={major_opt}, psum={use_psum_layout}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s') + print() + + +def test_m_grouped_gemm_masked() -> None: + print('Testing m-grouped masked GEMM:') + + # TODO: when the actual `m` is greater than `expected_m_per_group`, efficiency may significantly decrease. + for kernel_type, quant_config, num_groups, max_m, expected_m_per_group, n, k, use_psum_layout in enumerate_m_grouped_masked(torch.float8_e4m3fn): + kernel_opt = f'1D1D' if kernel_type.is_1d1d() else '1D2D' + use_ue8m0 = get_ue8m0_usage(kernel_type) + disable_ue8m0_cast = not use_ue8m0 + recipe, recipe_a, recipe_b = quant_config.get_recipes() + + num_tests = 8 + sum_t, max_t = 0, 0 + sum_ops, sum_bytes = 0, 0 + + # Select best alignment + alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(int(expected_m_per_group * 1.2)) + deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) + + for i in range(num_tests): + a, b, masked_m, psum_m, d, ref_d = generate_m_grouped_masked(num_groups, max_m, expected_m_per_group, n, k, + use_ue8m0=use_ue8m0, use_psum_layout=use_psum_layout, + quant_config=quant_config) + if use_psum_layout: + a_psum = (layout_masked_to_psum(a[0], psum_m), layout_masked_to_psum(a[1], psum_m)) + d_psum = layout_masked_to_psum(d, psum_m) + + # noinspection PyShadowingNames + def test_func(): + if use_psum_layout: + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(a_psum, b, d_psum, psum_m, disable_ue8m0_cast=disable_ue8m0_cast, + use_psum_layout=True, expected_m_for_psum_layout=int(expected_m_per_group * 1.2), + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + else: + deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, int(expected_m_per_group * 1.2), disable_ue8m0_cast=disable_ue8m0_cast, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + + test_func() + for j in range(num_groups): + if masked_m[j].item() == 0: + continue + if use_psum_layout: + d_slice = d_psum[: psum_m[j]] if j == 0 else d_psum[align(psum_m[j - 1], get_mk_alignment_for_contiguous_layout()): psum_m[j]] + else: + d_slice = d[j, :masked_m[j].item()] + diff = calc_diff(d_slice, ref_d[j, :masked_m[j].item()]) + assert diff < quant_config.max_diff(), f'{max_m=}, {n=}, {k=}, {j=}, masked_m={masked_m[j]}, {kernel_opt}, {num_groups=}, {diff:.5f}' + + # Test performance with fixed shapes + valid_m = masked_m.sum().item() + t = bench_kineto(test_func, 'gemm_', suppress_kineto_output=True) + + sum_t += t + max_t = max(max_t, t) + sum_ops += 2 * valid_m * n * k + sum_bytes += count_bytes(a, d) * valid_m / (max_m * num_groups) + count_bytes(b) + + print(f' > Perf (num_groups={num_groups:2}, expected_m_per_group={expected_m_per_group:4}, n={n:4}, k={k:4}, ' + f'{kernel_opt}, psum={1 if use_psum_layout else 0}): ' + f'{sum_t / num_tests * 1e6:4.0f} us (max: {max_t * 1e6:3.0f} us) | ' + f'{sum_ops / sum_t / 1e12:4.0f} TFLOPS | ' + f'{sum_bytes / sum_t / 1e9:4.0f} GB/s') + print() + + +def test_k_grouped_gemm_contiguous() -> None: + print('Testing k-grouped contiguous GEMM:') + + k_grouped_fp8_gemm_contiguous = deep_gemm.k_grouped_fp8_gemm_nt_contiguous if get_arch_major() == 9 \ + else deep_gemm.k_grouped_fp8_gemm_tn_contiguous + for num_groups, m, n, major_a, major_b, ks, expected_k_per_group, gran_k in enumerate_k_grouped_contiguous(torch.float8_e4m3fn): + recipe = (1, 1, gran_k) + use_ue8m0 = get_ue8m0_usage(KernelType.Kernel1D1D) + + for test_empty_groups in (False, True): + new_ks = copy.deepcopy(ks) + if test_empty_groups and len(ks) > 1: + new_ks[random.randint(0, num_groups - 1)] = 0 + k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, new_ks, use_ue8m0=use_ue8m0, gran_k=gran_k) + new_ks_tensor = torch.tensor(new_ks, dtype=torch.int, device='cuda') + k_grouped_fp8_gemm_contiguous(a, b, d, new_ks, new_ks_tensor, c, recipe=recipe) + + diff = calc_diff(d, ref_d) + assert diff < 0.001, f'{m=}, {n=}, {k=}, {ks=}, {diff:.5f}' + + # Test performance + k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_ue8m0=use_ue8m0, gran_k=gran_k) + ks_tensor = torch.tensor(ks, dtype=torch.int, device='cuda') + + # noinspection PyShadowingNames + def test_func(): + k_grouped_fp8_gemm_contiguous(a, b, d, ks, ks_tensor, c, recipe=recipe) + + t = bench_kineto(test_func, 'gemm_', suppress_kineto_output=True) + print(f' > Perf ({num_groups=:2}, m={m:5}, n={n:5}, k={k:5}, gran_k={gran_k:3}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, c, d) / 1e9 / t:4.0f} GB/s') + print() + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + + print('Library path:') + print(f' > {deep_gemm.__path__}\n') + + # test_gemm() + test_m_grouped_gemm_contiguous() + # test_m_grouped_gemm_masked() + # test_k_grouped_gemm_contiguous() diff --git a/sgl_deep_gemm/tests/test_hyperconnection.py b/sgl_deep_gemm/tests/test_hyperconnection.py new file mode 100644 index 0000000000..24faf22cbb --- /dev/null +++ b/sgl_deep_gemm/tests/test_hyperconnection.py @@ -0,0 +1,57 @@ +import torch +import random + +import deep_gemm +from deep_gemm.testing import ( + test_filter, + bench_kineto, + calc_diff, count_bytes +) +from deep_gemm.utils import align +from generators import get_arch_major + + +@test_filter(lambda: get_arch_major() >= 9) +def test_hc_prenorm_gemm() -> None: + # Needs TF32 precision for PyTorch GEMMs + torch.backends.cuda.matmul.allow_tf32 = True + torch.backends.cudnn.allow_tf32 = True + + print('Testing hyperconnection prenorm GEMM:') + for m in (13, 137, 4096, 8192): + for n, k in [(24, 28672), (24, 7680), (24, 7168)]: + for num_splits in [None, 16]: + a = torch.randn((m, k), dtype=torch.bfloat16, device='cuda') + b = torch.randn((n, k), dtype=torch.float, device='cuda') + d = torch.empty((m, n), dtype=torch.float, device='cuda') if num_splits is None else \ + torch.empty((num_splits, m, n), dtype=torch.float, device='cuda') + s = torch.empty((m, ), dtype=torch.float, device='cuda') if num_splits is None else \ + torch.empty((num_splits, m), dtype=torch.float, device='cuda') + deep_gemm.tf32_hc_prenorm_gemm(a, b, d, s, num_splits=num_splits) + final_d = d if num_splits is None else d.sum(0) + final_s = s if num_splits is None else s.sum(0) + + ref_d = a.float() @ b.T + ref_s = a.float().square().sum(-1) + + diff = max(calc_diff(final_d, ref_d), calc_diff(final_s, ref_s)) + assert diff < 1e-8, f'{m=}, {n=}, {k=}, {diff:.10f}' + + t = bench_kineto(lambda: deep_gemm.tf32_hc_prenorm_gemm(a, b, d, s, num_splits=num_splits), 'tf32_hc_prenorm_gemm', suppress_kineto_output=True) + print(f' > Perf (m={m:5}, n={n:5}, k={k:5}, num_splits={(num_splits or 0):2}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, d, s) / 1e9 / t:4.0f} GB/s') + print() + + + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + + print('Library path:') + print(f' > {deep_gemm.__path__}\n') + + test_hc_prenorm_gemm() diff --git a/sgl_deep_gemm/tests/test_layout.py b/sgl_deep_gemm/tests/test_layout.py new file mode 100644 index 0000000000..a0d4a02ebd --- /dev/null +++ b/sgl_deep_gemm/tests/test_layout.py @@ -0,0 +1,112 @@ +import torch +import random +from deep_gemm.testing import bench_kineto, count_bytes, get_arch_major +from deep_gemm.utils import ( + align, ceil_div, + per_token_cast_to_fp8, per_channel_cast_to_fp8, + get_tma_aligned_size, + get_mn_major_tma_aligned_tensor, + get_mn_major_tma_aligned_packed_ue8m0_tensor, + get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor +) + +from generators import ( + enumerate_sf_layout, + enumerate_k_grouped_sf_layout +) + + +def get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(x: torch.Tensor) -> torch.Tensor: + assert x.dtype == torch.float and x.dim() in (2, 3) + + # First, convert into UE8M0 `uint8_t` + ue8m0_tensor = (x.view(torch.int) >> 23).to(torch.uint8) + + # Second, make padded packed tensors + mn, k = x.shape[-2], x.shape[-1] + remove_dim = False + if x.dim() == 2: + x, remove_dim = x.unsqueeze(0), True + b = x.shape[0] + aligned_mn = get_tma_aligned_size(mn, 4) + aligned_k = align(k, 4) + padded = torch.zeros((b, aligned_mn, aligned_k), device=x.device, dtype=torch.uint8) + padded[:, :mn, :k] = ue8m0_tensor + padded = padded.view(-1).view(dtype=torch.int).view(b, aligned_mn, aligned_k // 4) + + # Finally, transpose + transposed = torch.zeros((b, aligned_k // 4, aligned_mn), device=x.device, dtype=torch.int).mT + transposed[:, :, :] = padded + aligned_x = transposed[:, :mn, :] + return aligned_x.squeeze(0) if remove_dim else aligned_x + + +def test_sf_layout_kernels() -> None: + print('Testing SF layout kernels:') + for mn, k, with_transpose, use_ue8m0, num_groups, gran_k in enumerate_sf_layout(): + x = torch.randn((num_groups * mn, k), dtype=torch.bfloat16, device='cuda') + x, fp32_sf = per_token_cast_to_fp8(x, use_ue8m0=use_ue8m0, gran_k=gran_k) + fp32_sf = fp32_sf if num_groups == 1 else fp32_sf.view(num_groups, mn, -1) + fp32_sf = fp32_sf if with_transpose else fp32_sf.transpose(-1, -2).contiguous().transpose(-1, -2) + + # Correctness + if use_ue8m0: + impl, name = get_mn_major_tma_aligned_packed_ue8m0_tensor, 'pack_fp32_into_ue8m0' + packed_sf = get_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf) + ref_packed_sf = get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(fp32_sf) + assert torch.equal(packed_sf, ref_packed_sf), f'{mn=}, {k=}, {with_transpose=}, {num_groups=}' + assert packed_sf.shape == ref_packed_sf.shape + assert all([packed_sf.stride(i) == ref_packed_sf.stride(i) for i in range(packed_sf.dim())]) + else: + impl, name = get_mn_major_tma_aligned_tensor, 'transpose' + transposed_sf = get_mn_major_tma_aligned_tensor(fp32_sf) + tma_aligned_mn, sf_k = get_tma_aligned_size(mn, fp32_sf.element_size()), ceil_div(k, gran_k) + if num_groups > 1: + assert transposed_sf.size(0) == num_groups + assert transposed_sf.stride(0) == tma_aligned_mn * sf_k + assert transposed_sf.shape[-2:] == (mn, sf_k) + assert transposed_sf.stride()[-2:] == (1, tma_aligned_mn) + assert torch.equal(fp32_sf, transposed_sf) + + # Performance + try: + t = bench_kineto(lambda: impl(fp32_sf), name) + except AssertionError as e: + # Some cases may fallback to PyTorch impl + t = 0 + print(f' > Perf ({num_groups=:2}, {mn=:5}, {k=:5}, transpose={int(with_transpose)}, use_ue8m0={int(use_ue8m0)}, gran_k={gran_k:3}): ' + f'{t * 1e6:4.0f} us | {count_bytes(fp32_sf, impl(fp32_sf)) / 1e9 / t if t else 0:4.0f} GB/s') + print() + + +def test_k_grouped_sf_layout_kernels() -> None: + print('Testing k-grouped SF layout kernels:') + for mn, ks, num_groups, gran_k in enumerate_k_grouped_sf_layout(): + sf_ks = [k // gran_k for k in ks] + packed_sf_ks = [ceil_div(k, gran_k * 4) for k in ks] + ks_tensor = torch.tensor(ks, dtype=torch.int, device='cuda') + x = torch.randn((sum(ks), mn), dtype=torch.bfloat16, device='cuda') + x, fp32_sf = per_channel_cast_to_fp8(x, use_ue8m0=True, gran_k=gran_k) + + # Correctness + packed_sf = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, ks_tensor, ks, gran_k) + split_packed_sf = packed_sf.split(packed_sf_ks) + split_fp32_sf = fp32_sf.split(sf_ks) + for i in range(num_groups): + ref_packed_sf = get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(split_fp32_sf[i].T).T + assert torch.equal(split_packed_sf[i], ref_packed_sf), f'{i=}' + + # Performance + t = bench_kineto(lambda: get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, ks_tensor, ks, gran_k), 'pack_fp32_into_ue8m0') + print(f' > Perf ({num_groups=:3}, {mn=:5}, sum_k={sum(ks):5}, gran_k={gran_k:3}):' + f'{t * 1e6:4.0f} us | ' + f'{count_bytes(fp32_sf, packed_sf, ks_tensor) / 1e9 / t:4.0f} GB/s') + print() + + +if __name__ == '__main__': + torch.manual_seed(1) + random.seed(1) + + test_sf_layout_kernels() + test_k_grouped_sf_layout_kernels() diff --git a/sgl_deep_gemm/tests/test_lazy_init.py b/sgl_deep_gemm/tests/test_lazy_init.py new file mode 100644 index 0000000000..17a3a121e4 --- /dev/null +++ b/sgl_deep_gemm/tests/test_lazy_init.py @@ -0,0 +1,20 @@ +import argparse +import torch +import torch.multiprocessing as mp +import deep_gemm + + +def main(local_rank: int): + torch.cuda.set_device(local_rank) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Test lazy initialization') + parser.add_argument('--num-processes', type=int, default=8, help='Number of processes to spawn (default: 8)') + args = parser.parse_args() + + procs = [mp.Process(target=main, args=(i, ), ) for i in range(args.num_processes)] + for p in procs: + p.start() + for p in procs: + p.join() diff --git a/sgl_deep_gemm/tests/test_legacy.py b/sgl_deep_gemm/tests/test_legacy.py new file mode 100644 index 0000000000..4456799f51 --- /dev/null +++ b/sgl_deep_gemm/tests/test_legacy.py @@ -0,0 +1,90 @@ +import torch +import random + +import deep_gemm +from deep_gemm.testing import ( + bench_kineto, + calc_diff, count_bytes +) +from generators import ( + enumerate_m_grouped_contiguous, enumerate_k_grouped_contiguous, + generate_m_grouped_contiguous, generate_k_grouped_contiguous, +) + +def test_m_grouped_gemm_contiguous_tl() -> None: + print('Testing m-grouped contiguous Triton GEMM:') + for _, _, num_groups, expected_m_per_group, n, k, major_a, major_b, _ in enumerate_m_grouped_contiguous(torch.bfloat16): + major_opt = 'N' if major_a.is_k_major() else 'T' + major_opt += 'T' if major_b.is_k_major() else 'N' + + for expand in (False, True): + for test_alias in (False, True): + m, a, b, m_indices, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, use_bf16=True) + func_name = f"{'a_fused_' if expand else ''}m_grouped_bf16_gemm_{major_opt.lower() if test_alias else 'nt'}_contiguous_tl" + if test_alias: + assert major_a.is_k_major() + b = b if major_b.is_k_major() else b.mT + assert a[0].is_contiguous() and b[0].is_contiguous() + if expand: + m_row_indices = torch.arange(0, m, dtype=torch.int32, device='cuda') + getattr(deep_gemm.legacy, func_name)(a, b, d, (m_indices, m_row_indices)) + else: + getattr(deep_gemm.legacy, func_name)(a, b, d, m_indices) + d = torch.where((m_indices == -1).unsqueeze(1), torch.zeros_like(d), d) + diff = calc_diff(d, ref_d) + assert diff < 0.001, f'{m=}, {n=}, {k=}, {major_opt}, {diff:.5f}, alias={test_alias}' + m, a, b, m_indices, d, ref_d = generate_m_grouped_contiguous(num_groups, expected_m_per_group, n, k, major_a, major_b, use_bf16=True) + + # noinspection PyShadowingNames + def test_func(): + deep_gemm.legacy.m_grouped_bf16_gemm_nt_contiguous_tl(a, b, d, m_indices) + + t = bench_kineto(test_func, 'm_grouped_bf16_gemm_contiguous_tl_impl', suppress_kineto_output=True) + print(f' > Perf ({num_groups=}, m={m:5}, n={n:5}, k={k:5}, layout={major_opt}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s') + print() + + +def test_k_grouped_gemm_contiguous_tl() -> None: + print('Testing k-grouped contiguous Triton GEMM:') + for num_groups, m, n, major_a, major_b, ks, expected_k_per_group in enumerate_k_grouped_contiguous(torch.bfloat16): + major_opt = 'N' if major_a.is_k_major() else 'T' + major_opt += 'T' if major_b.is_k_major() else 'N' + + for fused_operand in ('a', 'b'): + k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_ue8m0=False, use_bf16=True) + func_name = f"{fused_operand}_fused_k_grouped_bf16_gemm_{major_opt.lower()}_contiguous_tl" + k_indices = torch.arange(0, k, dtype=torch.int32, device='cuda') + k_start = torch.empty(len(ks), dtype=torch.int32, device='cuda') + k_end = torch.empty(len(ks), dtype=torch.int32, device='cuda') + for i, group_k in enumerate(ks): + k_start[i] = k_end[i-1] if i > 0 else 0 + k_end[i] = k_start[i] + group_k + getattr(deep_gemm.legacy, func_name)(a, b, c, (k_indices, k_start, k_end), True) + diff = calc_diff(c, ref_d) + assert diff < 0.001, f'{m=}, {n=}, {k=}, {major_opt}, {diff:.5f}' + k, a, b, c, d, ref_d = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, ks, use_ue8m0=False, use_bf16=True) + + # noinspection PyShadowingNames + def test_func(): + deep_gemm.legacy.b_fused_k_grouped_bf16_gemm_tn_contiguous_tl(a, b, c, (k_indices, k_start, k_end), True) + + t = bench_kineto(test_func, 'b_fused_k_grouped_bf16_gemm_contiguous_tl_impl', suppress_kineto_output=True) + print(f' > Perf ({num_groups=}, m={m:5}, n={n:5}, k={k:5}, layout={major_opt}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, d) / 1e9 / t:4.0f} GB/s') + print() + + +if __name__ == '__main__': + torch.manual_seed(0) + random.seed(0) + + print('Library path:') + print(f' > {deep_gemm.__path__}\n') + + test_m_grouped_gemm_contiguous_tl() + test_k_grouped_gemm_contiguous_tl() diff --git a/sgl_deep_gemm/tests/test_mega_moe.py b/sgl_deep_gemm/tests/test_mega_moe.py new file mode 100644 index 0000000000..5111edda23 --- /dev/null +++ b/sgl_deep_gemm/tests/test_mega_moe.py @@ -0,0 +1,301 @@ +import argparse +import os +import random +import sys +import torch +import torch.distributed as dist +from typing import Tuple + +import deep_gemm +from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8 +from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather +from deep_gemm.testing import bench_kineto + + +def import_baseline(): + # Load legacy implements from third-party + deep_ep, tilelang_ops, do_bench, is_legacy_loaded = None, None, None, False + # noinspection PyBroadException + try: + import deep_ep + import importlib.util + from tilelang.profiler.bench import do_bench + spec = importlib.util.spec_from_file_location( + 'tilelang_ops', + os.path.join(os.path.dirname(os.path.realpath(__file__)), '..', 'third-party', 'tilelang_ops', '__init__.py')) + tilelang_ops = importlib.util.module_from_spec(spec) + sys.modules['tilelang_ops'] = tilelang_ops + spec.loader.exec_module(tilelang_ops) + is_legacy_loaded = True + except Exception as ex: + dist_print(f'Failed to load legacy code: {ex}, skip baseline benchmarking', once_in_node=True) + dist_print(once_in_node=True) + return deep_ep, tilelang_ops, do_bench, is_legacy_loaded + + +# TODO: skip the test for SM90 +# noinspection PyUnboundLocalVariable,PyShadowingNames +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + torch.manual_seed(rank_idx) + random.seed(rank_idx) + + # Settings + num_max_tokens_per_rank = args.num_max_tokens_per_rank + num_tokens = max(0, args.num_max_tokens_per_rank - random.randint(0, args.num_max_removed_tokens)) \ + if args.num_tokens == 0 else args.num_tokens + hidden, intermediate_hidden = args.hidden, args.intermediate_hidden + num_experts, num_topk = args.num_experts, args.num_topk + num_experts_per_rank = num_experts // num_ranks + assert num_tokens <= num_max_tokens_per_rank + + # Allocate symmetric memory + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden + ) + + # Create inputs + # noinspection PyGlobalUndefined + def create_inputs(): + global x, topk_idx, topk_weights, l1_weights, l2_weights, transformed_l1_weights, transformed_l2_weights + global cumulative_local_expert_recv_stats_fused + global cumulative_local_expert_recv_stats_baseline + x = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + l1_weights = torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), dtype=torch.bfloat16, device='cuda') + l2_weights = torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), dtype=torch.bfloat16, device='cuda') + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device='cuda') + topk_weights, topk_idx = torch.topk(scores, num_topk, dim=-1, largest=True, sorted=False) + cumulative_local_expert_recv_stats_fused = torch.randint( + 0, 100, (num_experts_per_rank, ), dtype=torch.int, device='cuda') + cumulative_local_expert_recv_stats_baseline = cumulative_local_expert_recv_stats_fused.clone() + if args.masked_ratio > 0: + rand_mask = torch.rand_like(topk_idx, dtype=torch.float) + topk_idx.masked_fill_(rand_mask < args.masked_ratio, -1) + topk_weights.masked_fill_(topk_idx < 0, 0) + + # Check SF requirements + assert hidden % 128 == 0 + assert intermediate_hidden % 128 == 0 + assert l1_weights.shape[2] % 128 == 0 and l2_weights.shape[2] % 128 == 0 + + # Cast inputs to FP8 (or FP4 under DG_USE_FP4_ACTS) with per-32 UE8M0 SF. + # Stream A0.0b: when the flag is on, the symm buffer's `x` slot is sized + # for packed E2M1 (`hidden/2` bytes/token), so we must quantize at the + # source to match. + if os.environ.get('DG_USE_FP4_ACTS', '0') != '0': + x = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + else: + x = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + + # Cast grouped BF16 weights to FP4 with MN-major SF + # TODO: merge with `cast_fp8_fp4_with_major` + def cast_grouped_weights_to_fp4(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + num_groups, n, k = bf16_weights.shape + w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) + w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) + for i in range(num_groups): + w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) + w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + return w, w_sf + + l1_weights = cast_grouped_weights_to_fp4(l1_weights) + l2_weights = cast_grouped_weights_to_fp4(l2_weights) + transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights) + + # Run fused mega MoE + # NOTES: copy x into buffer before each call because debug mode zeros the entire buffer + def run_fused(): + buffer.x[:num_tokens].copy_(x[0]) + buffer.x_sf[:num_tokens].copy_(x[1]) + buffer.topk_idx[:num_tokens].copy_(topk_idx) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + + y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') + # noinspection PyTypeChecker + deep_gemm.fp8_fp4_mega_moe( + y, + transformed_l1_weights, transformed_l2_weights, + buffer, + cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused, + activation_clamp=args.activation_clamp, + fast_math=bool(args.fast_math) + ) + return y, cumulative_local_expert_recv_stats_fused + + dist_print('Config:', once_in_node=True) + dist_print(f' > Tokens: {num_tokens}/{num_max_tokens_per_rank}', once_in_node=True) + dist_print(f' > Hidden: {hidden}', once_in_node=True) + dist_print(f' > Intermediate: {intermediate_hidden}', once_in_node=True) + dist_print(f' > Experts: {num_topk}/{num_experts}', once_in_node=True) + dist_print(f' > Buffer: {buffer.buffer.nbytes / 2 ** 30:.3f} GiB', once_in_node=True) + dist_print(once_in_node=True) + + # Only do NCU profiling + if args.ncu_profile_only: + create_inputs() + dist_print(f'Run fused kernel:', once_in_node=True) + run_fused() + dist_print(f' > Done, exiting', once_in_node=True) + + # Destroy and exit + dist.barrier() + buffer.destroy() + dist.destroy_process_group() + return + + # Non-overlapped baseline: EP dispatch + GEMM + EP combine + deep_ep, tilelang_ops, tilelang_bench, is_legacy_loaded = import_baseline() + alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout() + deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) + ep_buffer = deep_ep.ElasticBuffer( + group, + num_max_tokens_per_rank=num_max_tokens_per_rank, hidden=hidden, + num_topk=num_topk, use_fp8_dispatch=True, + explicitly_destroy=True, + allow_multiple_reduction=False, + gpu_timeout_secs=10, cpu_timeout_secs=30 + ) if is_legacy_loaded else None + + def run_baseline(): + recv_x, _, recv_topk_weights, handle, _ = ep_buffer.dispatch( + x, topk_idx=topk_idx, topk_weights=topk_weights, + cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_baseline, + num_experts=num_experts, expert_alignment=alignment, + do_cpu_sync=False, do_handle_copy=False, + do_expand=True, use_tma_aligned_col_major_sf=True, + ) + n = recv_x[0].size(0) + l1_y = torch.empty((n, intermediate_hidden * 2), dtype=torch.bfloat16, device='cuda') + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( + recv_x, l1_weights, l1_y, handle.psum_num_recv_tokens_per_expert, + use_psum_layout=True, recipe=(1, 1, 32)) + # noinspection PyCallingNonCallable + l1_y = tilelang_ops.swiglu_apply_weight_to_fp8( + x=l1_y, + topk_weights=recv_topk_weights, + avail_tokens=handle.psum_num_recv_tokens_per_expert[-1], + num_per_channels=32, + use_col_major_scales=True, + round_scale=True, + ue8m0_scale=True, + output_bf16=False, + clamp_value=args.activation_clamp, + fast_math=bool(args.fast_math) + ) + l2_y = torch.empty((n, hidden), dtype=torch.bfloat16, device='cuda') + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( + l1_y, l2_weights, l2_y, handle.psum_num_recv_tokens_per_expert, + use_psum_layout=True, recipe=(1, 1, 32)) + return ep_buffer.combine(l2_y, handle=handle)[0], cumulative_local_expert_recv_stats_baseline + + # Check correctness (must be bitwise identical) + num_correctness_tests = 1 if args.num_correctness_tests is None else args.num_correctness_tests + # noinspection PyBroadException + if is_legacy_loaded and num_correctness_tests > 0: + dist_print('Running correctness tests:', once_in_node=True) + for i in range(num_correctness_tests): + create_inputs() + for fused_result, baseline_result in zip(run_fused(), run_baseline()): + assert torch.equal(fused_result, baseline_result) + if (i + 1) % 100 == 0 or i == num_correctness_tests - 1: + dist_print(f' > Correctness test #{i + 1}/{num_correctness_tests} passed', once_in_node=True) + dist_print(once_in_node=True) + else: + create_inputs() + + # Count local received tokens + gathered_topk_idx = uneven_all_gather(topk_idx, group=group) + gathered_topk_idx[(gathered_topk_idx < rank_idx * num_experts_per_rank) | \ + (gathered_topk_idx >= (rank_idx + 1) * num_experts_per_rank)] = -1 + num_recv_tokens = (gathered_topk_idx != -1).sum().item() + + # Benchmark + t_fused = bench_kineto( + run_fused, 'mega_moe', + barrier=lambda: ep_buffer.barrier(use_comm_stream=False) if ep_buffer else dist.barrier(), + trace_path=None if not args.dump_profile_traces else f'{args.dump_profile_traces}/mega_moe_rank{rank_idx}.json') + t_baseline = tilelang_bench(run_baseline, _n_warmup=5, _n_repeat=1, backend='cudagraph', return_mode='median') / 1e3 if is_legacy_loaded else 0 + + # TFLOPS: 3 matmuls (L1 left, L1 right, L2), each 2 * M * N * K + safe_div = lambda a, b: float('nan') if b == 0 else a / b + tflops = safe_div(2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_fused) + + # HBM bytes: weights (FP4 packed = 0.5 bytes) + activations (FP8 = 1 byte) + output (BF16 = 2 bytes) + num_touched_experts = torch.unique(gathered_topk_idx.flatten()).numel() - 1 # NOTES minus 1 to exclude "-1" + num_hbm_bytes = ( + num_touched_experts * intermediate_hidden * 2 * hidden // 2 + # L1 weights (FP4) + num_touched_experts * hidden * intermediate_hidden // 2 + # L2 weights (FP4) + num_recv_tokens * hidden + # L1 acts read (FP8) + num_recv_tokens * intermediate_hidden + # L1 output write (FP8) + num_recv_tokens * intermediate_hidden + # L2 acts read (FP8) + num_recv_tokens * hidden * 2 # L2 output write (BF16) + ) + hbm_gbs = safe_div(num_hbm_bytes / 1e9, t_fused) + + # NVLink bytes: dispatch pull + combine write-back + num_nvlink_bytes = num_recv_tokens * hidden * 3 + nvlink_gbs = safe_div(num_nvlink_bytes / 1e9, t_fused) + + # Combine reduction (serial) time approximation + t_reduction = num_tokens * hidden * 2 * (1 + num_topk) / 6.5e12 + + # Summary + approx_factor = t_fused / (t_fused - t_reduction) + dist_print('Performance:', once_in_node=True) + dist_print(f' > EP: {rank_idx:2}/{num_ranks} | ' + f'{tflops:4.0f} TFLOPS | ' + f'overlap: ' + f'{tflops * approx_factor:4.0f} TFLOPS, ' + f'HBM {hbm_gbs * approx_factor:4.0f} GB/s, ' + f'NVL {nvlink_gbs * approx_factor:3.0f} GB/s | ' + f'{t_fused * 1e6:4.0f} us, ' + f'reduction: {t_reduction * 1e6:4.1f} us | ' + f'{safe_div(t_baseline, t_fused):.2f}x legacy') + + # Exit + dist.barrier() + buffer.destroy() + ep_buffer.destroy() if is_legacy_loaded else None + dist.destroy_process_group() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Test PyTorch symmetric memory') + + # Resource settings + parser.add_argument('--ncu-profile-only', action='store_true', help='Only run profiling without correctness test') + parser.add_argument('--num-processes', type=int, default=8, help='Number of processes to spawn (default: 8)') + + # Model settings + parser.add_argument('--num-max-tokens-per-rank', type=int, default=8192, help='Number of maximum tokens per rank') + parser.add_argument('--num-tokens', type=int, default=0, help='Number of tokens per rank (follow max minus removed if 0)') + parser.add_argument('--num-max-removed-tokens', type=int, default=0, help='Maximum number of tokens to remove') + parser.add_argument('--hidden', type=int, default=7168, help='Hidden size') + parser.add_argument('--intermediate-hidden', type=int, default=3072, help='Intermediate hidden size') + parser.add_argument('--activation-clamp', type=float, default=10, help='Clamp value for activation') + parser.add_argument('--num-experts', type=int, default=384, help='Number of experts') + parser.add_argument('--num-topk', type=int, default=6, help='Number of expert selections') + parser.add_argument('--masked-ratio', type=float, default=0.0, help='Mask some expert selections') + parser.add_argument('--fast-math', type=int, default=1, help='Enable fast math (0 or 1, default: 1)') + + # Test settings + parser.add_argument('--num-correctness-tests', type=int, default=None, help='Pressure test') + parser.add_argument('--dump-profile-traces', type=str, default='', help='Dump profiling trace JSONs') + parser.add_argument('--local-rank-idx', type=int, default=None, help='Run as single process with this local rank (e.g. for NCU prof)') + args = parser.parse_args() + + # Create dump trace directories + if args.dump_profile_traces: + os.makedirs(args.dump_profile_traces, exist_ok=True) + + if args.local_rank_idx is not None: + # Single-process mode: each process is launched separately (e.g. by NCU) + test(args.local_rank_idx, args.num_processes, args) + else: + # Launch tests + num_processes = args.num_processes + torch.multiprocessing.spawn(test, args=(num_processes, args), nprocs=num_processes) diff --git a/tests/test_mega_moe_l1_fp4_accuracy.py b/sgl_deep_gemm/tests/test_mega_moe_l1_fp4_accuracy.py similarity index 95% rename from tests/test_mega_moe_l1_fp4_accuracy.py rename to sgl_deep_gemm/tests/test_mega_moe_l1_fp4_accuracy.py index 1d013ff240..1cd77fba6f 100644 --- a/tests/test_mega_moe_l1_fp4_accuracy.py +++ b/sgl_deep_gemm/tests/test_mega_moe_l1_fp4_accuracy.py @@ -241,12 +241,6 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): activation_clamp = args.activation_clamp assert num_tokens <= num_max_tokens_per_rank - buffer = deep_gemm.get_symm_buffer_for_mega_moe( - group, num_experts, - num_max_tokens_per_rank, num_topk, - hidden, intermediate_hidden - ) - # Inputs (BF16) + topk routing x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') l1_weights_bf16 = torch.randn( @@ -262,6 +256,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): # FP8 / FP4 quantizations needed by the kernel x_fp8 = per_token_cast_to_fp8(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + x_fp4 = per_token_cast_to_fp4(x_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) def cast_grouped_weights_to_fp4(bf16_weights): num_groups, n, k = bf16_weights.shape @@ -277,9 +272,9 @@ def cast_grouped_weights_to_fp4(bf16_weights): transformed_l1_weights, transformed_l2_weights = \ deep_gemm.transform_weights_for_mega_moe(l1_weights_fp4, l2_weights_fp4) - def run_once(): - buffer.x[:num_tokens].copy_(x_fp8[0]) - buffer.x_sf[:num_tokens].copy_(x_fp8[1]) + def run_once(buffer, x_src): + buffer.x[:num_tokens].copy_(x_src[0]) + buffer.x_sf[:num_tokens].copy_(x_src[1]) buffer.topk_idx[:num_tokens].copy_(topk_idx) buffer.topk_weights[:num_tokens].copy_(topk_weights) cumulative_local_expert_recv_stats.zero_() @@ -294,27 +289,31 @@ def run_once(): ) return y, cumulative_local_expert_recv_stats.clone() + # Buffer layout depends on DG_USE_FP4_ACTS; set the env before allocating. + def make_buffer(use_fp4_acts): + os.environ['DG_USE_FP4_ACTS'] = '1' if use_fp4_acts else '0' + os.environ['DG_COMM_KERNEL_DEBUG'] = '0' # don't zero buffer between calls + return deep_gemm.get_symm_buffer_for_mega_moe( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden + ) + # ---- BF16 reference for L1 SwiGLU output (per token×topk) ---- bf16_ref = _bf16_reference_l1( x_bf16, l1_weights_bf16, topk_idx, topk_weights, activation_clamp) # bf16_ref: (num_tokens, num_topk, intermediate_hidden) — only nonzero # where topk_idx[t, k] is in this rank's expert range. - # ---- Run FP8 path ---- - os.environ['DG_USE_FP4_ACTS'] = '0' - os.environ['DG_COMM_KERNEL_DEBUG'] = '0' # don't zero buffer between calls - # First run is a warmup. Stream A0.2 verified FP8-vs-FP8 across two - # consecutive runs gives a perfect (rel-MAE = 0) `y` match — the kernel - # IS deterministic at the `y` level, so any nonzero FP4-vs-FP8 `y` - # delta is a real numerical disagreement, not slot-permutation noise. - _ = run_once() + # ---- Run FP8 path (first call warms up) ---- + buffer_fp8 = make_buffer(use_fp4_acts=False) + _ = run_once(buffer_fp8, x_fp8) torch.cuda.synchronize() - y_fp8, recv_stats_fp8 = run_once() + y_fp8, recv_stats_fp8 = run_once(buffer_fp8, x_fp8) torch.cuda.synchronize() - y_fp8_a = y_fp8 # keep as alias so the FP8-vs-FP8 baseline below works - # Snapshot l2_acts and l2_acts_sf before they get overwritten by next call. - l2_acts_fp8 = buffer.l2_acts.clone() - l2_acts_sf_fp8 = buffer.l2_acts_sf.clone() + y_fp8_a = y_fp8 + l2_acts_fp8 = buffer_fp8.l2_acts.clone() + l2_acts_sf_fp8 = buffer_fp8.l2_acts_sf.clone() recv_fp8_list = recv_stats_fp8.cpu().tolist() # `recv_stats` is per-expert cumulative — the last element is the running # total of tokens routed to this rank's experts (since dispatcher @@ -322,11 +321,11 @@ def run_once(): # take the last value as the slot count. total_local_fp8 = int(recv_fp8_list[-1]) if recv_fp8_list else 0 - # ---- Run FP4 path ---- - os.environ['DG_USE_FP4_ACTS'] = '1' - _ = run_once() + # ---- Run FP4 path (separate buffer, laid out for packed E2M1) ---- + buffer = make_buffer(use_fp4_acts=True) + _ = run_once(buffer, x_fp4) torch.cuda.synchronize() - y_fp4, recv_stats_fp4 = run_once() + y_fp4, recv_stats_fp4 = run_once(buffer, x_fp4) torch.cuda.synchronize() l2_acts_fp4 = buffer.l2_acts.clone() l2_acts_sf_fp4 = buffer.l2_acts_sf.clone() diff --git a/tests/test_mega_moe_l1_sentinel.py b/sgl_deep_gemm/tests/test_mega_moe_l1_sentinel.py similarity index 100% rename from tests/test_mega_moe_l1_sentinel.py rename to sgl_deep_gemm/tests/test_mega_moe_l1_sentinel.py diff --git a/tests/test_mega_moe_pre_dispatch.py b/sgl_deep_gemm/tests/test_mega_moe_pre_dispatch.py similarity index 100% rename from tests/test_mega_moe_pre_dispatch.py rename to sgl_deep_gemm/tests/test_mega_moe_pre_dispatch.py diff --git a/sgl_deep_gemm/tests/test_sanitizer.py b/sgl_deep_gemm/tests/test_sanitizer.py new file mode 100644 index 0000000000..75ab10e60a --- /dev/null +++ b/sgl_deep_gemm/tests/test_sanitizer.py @@ -0,0 +1,79 @@ +import argparse +import importlib +import inspect +import os +import subprocess +import sys + +import deep_gemm + + +# Single test template +script_dir = os.path.dirname(os.path.abspath(__file__)) +test_template = """ +import random +import sys +import torch + +# Necessary for `generators.py` +sys.path.append('{script_dir}') + +torch.manual_seed(0) +random.seed(0) + +from {module_name} import {func_name} +{func_name}() +""" + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('--funcs', type=str, default='all') + parser.add_argument('--tools', type=str, default='memcheck,synccheck') + args = parser.parse_args() + + if args.funcs != 'all': + funcs = [] + for name in [x.strip() for x in args.funcs.split(',')]: + module_name, func_name = name.split('.') + funcs.append((module_name, func_name)) + else: + # Get all test functions except those related to cuBLAS + files = [f for f in os.listdir(script_dir) if f.endswith('.py')] + exclude_files = ['test_sanitizer.py', 'generators.py', 'test_mega_moe.py'] + funcs = [ + (module_name, name) + for module_name in [os.path.splitext(f)[0] for f in files if f not in exclude_files] + for name, obj in inspect.getmembers(importlib.import_module(module_name)) + if inspect.isfunction(obj) and name.startswith('test') and 'test_filter' not in name + ] + tools = [x.strip() for x in args.tools.split(',')] + + env = os.environ.copy() + env['CUDA_LAUNCH_BLOCKING'] = '1' + env['DG_JIT_PTXAS_CHECK'] = '1' + env['DG_USE_NVIDIA_TOOLS'] = '1' + env['DG_USE_TEMP_CUBLASLT_WORKSPACE'] = '1' # Avoid holding CUDA tensor that crashes during shutdown + env['PYTORCH_NO_CUDA_MEMORY_CACHING'] = '1' + env['TORCH_SHOW_CPP_STACKTRACES'] = '1' + + print(f'Library path: {deep_gemm.__path__}') + for module_name, func_name in funcs: + for tool in tools: + cmd = [ + '/usr/local/cuda/bin/compute-sanitizer', + f'--tool={tool}', + '--target-processes=application-only', + '--destroy-on-device-error=context', + '--force-blocking-launches', + '--check-api-memory-access=no', + '--kernel-name-exclude', 'kns=nvjet', + 'python', + '-c', + test_template.format(module_name=module_name, func_name=func_name, script_dir=script_dir) + ] + print(f'\n{"=" * 60}') + print(f'Running {module_name}.{func_name} with compute-sanitizer {tool}') + result = subprocess.run(cmd, env=env) + if result.returncode != 0: + sys.exit(result.returncode) diff --git a/tests/generators.py b/tests/generators.py index 989e984e7c..4ec09aff67 100644 --- a/tests/generators.py +++ b/tests/generators.py @@ -136,6 +136,9 @@ def enumerate_normal(dtype: torch.dtype) -> Generator: n, k = nk_list[i] out_dtype = torch.bfloat16 if i < len(bf16_output_nk) else torch.float yield kernel_type, quant_config, m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, False, out_dtype + # BF16 accumulation: supported on all BF16 GEMMs, and SM100 FP8/FP4 GEMMs + if out_dtype == torch.bfloat16 and (dtype == torch.bfloat16 or get_arch_major() == 10): + yield kernel_type, quant_config, m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, True, out_dtype # Backward for m in m_bwd_list: diff --git a/tests/test_attention.py b/tests/test_attention.py index 479da5b56f..769923a5c8 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -113,7 +113,7 @@ def enumerate_mqa_logits(): for compressed_logits, clean_logits in [(False, True), (True, False)]: for seq_len in (2048, 4096): for seq_len_kv in (4096, 8192): - for num_heads, head_dim in [(64, 128)]: + for num_heads, head_dim in [(64, 128), (32, 128)]: for disable_cp in (False, True): yield is_fp4, logits_dtype, compressed_logits, clean_logits, seq_len, seq_len_kv, num_heads, head_dim, disable_cp @@ -224,6 +224,18 @@ def ref_paged_mqa_logits(q: torch.Tensor, kv_cache: torch.Tensor, return logits +def reset_cuda_peak_memory_stats(): + torch.cuda.synchronize() + torch.cuda.reset_peak_memory_stats() + + +def get_cuda_peak_memory_gib(): + torch.cuda.synchronize() + peak_allocated = torch.cuda.max_memory_allocated() / 1024 ** 3 + peak_reserved = torch.cuda.max_memory_reserved() / 1024 ** 3 + return peak_allocated, peak_reserved + + def test_paged_mqa_logits(): # Helper functions @@ -253,24 +265,30 @@ def kv_cache_cast_to_fp4(x: torch.Tensor) -> torch.Tensor: def enumerate_paged_mqa_logits(): arch_major = get_arch_major() + max_kv_pool_tokens = 32 * 1024 * 1024 + max_varlen_tokens = 16 * 1024 for is_varlen in ((True, False) if arch_major == 10 else (False, )): for is_fp4 in ((True, False) if arch_major == 10 else (False, )): for logits_dtype in (torch.float, torch.bfloat16): for block_kv in ((32, 64) if arch_major == 10 else (64, )): for use_2d_context_lens, clean_logits in [(True, False)]: - for batch_size in (256, ): + for batch_size in (256, 4096): for next_n in ((1, ) if is_varlen else ((1, 2, 4, 5, 6) if arch_major == 10 else (1, 2))): for max_tokens_per_batch in ((1, 4, 10) if is_varlen else (1, )): - for num_heads, head_dim in [(64, 128)]: + for num_heads, head_dim in [(64, 128), (32, 128)]: for avg_kv in (8192, 32768): + if batch_size * avg_kv > max_kv_pool_tokens: + continue + if is_varlen and batch_size * max_tokens_per_batch > max_varlen_tokens: + continue yield is_varlen, is_fp4, logits_dtype, block_kv, use_2d_context_lens, clean_logits, batch_size, next_n, max_tokens_per_batch, num_heads, head_dim, avg_kv print('Testing FP8/FP4 Paged MQA Logits:') - max_model_len = 111 * 1024 - num_total_blocks = max_model_len * 5 for is_varlen, is_fp4, logits_dtype, block_kv, use_2d_context_lens, clean_logits, batch_size, next_n, max_tokens_per_batch, num_heads, head_dim, avg_kv in enumerate_paged_mqa_logits(): + reset_cuda_peak_memory_stats() + # Varlen: flatten raw_batch_size sequences with variable tokens into (batch_size, 1, ...) raw_batch_size, raw_next_n = batch_size, next_n if is_varlen: @@ -282,7 +300,6 @@ def enumerate_paged_mqa_logits(): # Generate random inputs q = torch.randn((batch_size, next_n, num_heads, head_dim), device='cuda', dtype=torch.bfloat16) - kv_cache = torch.randn((num_total_blocks, block_kv, 1, head_dim), device='cuda', dtype=torch.bfloat16) weights = torch.randn((batch_size * next_n, num_heads), device='cuda', dtype=torch.float) context_lens = torch.randint(int(0.7 * avg_kv), int(1.3 * avg_kv), (raw_batch_size,), device='cuda', dtype=torch.int) @@ -294,7 +311,10 @@ def enumerate_paged_mqa_logits(): # Assign block tables (per-sequence, sized by the largest ctx_len within the sequence) seq_sum_lens = context_lens.sum().item() num_blocks_per_query = ceil_div(max_ctx_len_per_seq, block_kv) - block_table = torch.empty((raw_batch_size, num_blocks_per_query.max().item()), device='cuda', dtype=torch.int) + max_model_len = num_blocks_per_query.max().item() * block_kv + num_total_blocks = num_blocks_per_query.sum().item() + kv_cache = torch.randn((num_total_blocks, block_kv, 1, head_dim), device='cuda', dtype=torch.bfloat16) + block_table = torch.zeros((raw_batch_size, num_blocks_per_query.max().item()), device='cuda', dtype=torch.int) block_idx_pool = torch.randperm(num_total_blocks, device='cuda', dtype=torch.int) offset = 0 for i, num_blocks in enumerate(num_blocks_per_query.tolist()): @@ -311,6 +331,7 @@ def enumerate_paged_mqa_logits(): # Calculate reference logits ref_logits = ref_paged_mqa_logits(q, kv_cache, weights, context_lens, block_table, max_model_len, use_2d_context_lens) + q_weight_bytes = count_bytes(q, weights) # Quantize Q and KV cache to FP4 / FP8 if is_fp4: @@ -322,6 +343,7 @@ def enumerate_paged_mqa_logits(): q_in = q.to(torch.float8_e4m3fn), None q_simulated = q_in[0].to(torch.bfloat16) kv_in, kv_simulated = kv_cache_cast_to_fp8(kv_cache) + del q, kv_cache # Calculate simulated reference logits simulated_logits = ref_paged_mqa_logits(q_simulated, kv_simulated, weights, context_lens, block_table, max_model_len, use_2d_context_lens) @@ -345,6 +367,9 @@ def enumerate_paged_mqa_logits(): ref_neginf_mask = ~(positions <= limits) # Run Kernel + assert block_table.min().item() >= 0 + assert block_table.max().item() < num_total_blocks + assert context_lens_nextn.max().item() <= max_model_len kernel_kwargs = dict( q=q_in, kv_cache=kv_in, weights=weights, context_lens=context_lens_nextn, block_table=block_table, @@ -375,14 +400,24 @@ def enumerate_paged_mqa_logits(): kv_bytes_per_token = head_dim / (2 if is_fp4 else 1) + 4 # KV is read once per sequence; for varlen sum_lens overcounts (per-token), so use seq_sum_lens kv_sum_lens = seq_sum_lens if is_varlen else sum_lens - total_bytes = count_bytes(q, weights) + kv_sum_lens * kv_bytes_per_token + (sum_lens * next_n * logits_dtype.itemsize) + total_bytes = q_weight_bytes + kv_sum_lens * kv_bytes_per_token + (sum_lens * next_n * logits_dtype.itemsize) + peak_allocated, peak_reserved = get_cuda_peak_memory_gib() t, clean_t = bench_kineto(lambda: deep_gemm.fp8_fp4_paged_mqa_logits(**kernel_kwargs), ('paged_mqa_logits', 'clean_logits')) print(f' > FP4={is_fp4}, BF16={logits_dtype == torch.bfloat16}, BLOCK_KV={block_kv}, BSZ={raw_batch_size:3}, NextN={raw_next_n:1}, H={num_heads:2}, D={head_dim:2}, L={avg_kv:6}: ' f'{tflops_calc / t:4.0f} TFLOPS, {t * 1e6:3.0f} us, {total_bytes / t / 1e9:4.0f} GB/s', end='') if is_varlen: print(f' | Varlen, MaxTPB={max_tokens_per_batch}, NumTokens={batch_size}', end='') + print(f' | mem: alloc={peak_allocated:.2f} GiB, reserved={peak_reserved:.2f} GiB', end='') print(f' | clean: {clean_t*1e6:3.0f} us' if clean_logits else '') + + del kernel_kwargs, logits, ref_neginf_mask, positions + del q_in, q_simulated, kv_in, kv_simulated, weights, context_lens, context_lens_nextn, block_table + if is_fp4: + del q_fp4 + if is_varlen: + del tokens_per_seq, indices, offsets_within_seq + torch.cuda.empty_cache() print() diff --git a/tests/test_bf16.py b/tests/test_bf16.py index 703337b86e..a3740e35d2 100644 --- a/tests/test_bf16.py +++ b/tests/test_bf16.py @@ -39,7 +39,7 @@ def test_gemm() -> None: a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_bf16=True) t = bench_kineto(lambda: deep_gemm.bf16_gemm_nt(a, b, d, c=c), 'bf16_gemm', suppress_kineto_output=True) - cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c), ('nvjet', 'reduce'), suppress_kineto_output=True) + cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c=c), ('nvjet', 'reduce'), suppress_kineto_output=True) print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, layout={major_opt}, {out_opt}, {acc_opt}): ' f'{t * 1e6:7.1f} us | ' f'{2 * m * n * k / t / 1e12:4.0f} TFLOPS | ' @@ -194,11 +194,13 @@ def test_cublaslt_gemm() -> None: acc_opt = f'acc={int(accumulate)}' a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_bf16=True) - deep_gemm.cublaslt_gemm_nt(a, b, d, c) + deep_gemm.cublaslt_gemm_nt(a, b, d, c=c) diff = calc_diff(d, ref_d) - assert diff < 6e-7, f'{diff=}, ({m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=})' + # BF16 accumulation has lower precision than cuBLASLt's FP32 accumulation + threshold = 1e-5 if (accumulate and out_dtype == torch.bfloat16) else 6e-7 + assert diff < threshold, f'{diff=}, ({m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=})' - t_nvjet, t_gemv, t_gemm = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c), ('nvjet', 'gemv', 'gemm'), suppress_kineto_output=True) + t_nvjet, t_gemv, t_gemm = bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c=c), ('nvjet', 'gemv', 'gemm'), suppress_kineto_output=True) t = t_nvjet + t_gemv + t_gemm print(f' > Perf (m={m:6}, n={n:6}, k={k:6}, layout={major_opt}, {out_opt}, {acc_opt}): ' f'{t * 1e6:5.0f} us | ' diff --git a/tests/test_fp8_fp4.py b/tests/test_fp8_fp4.py index 42038e459d..4e9f54f7f4 100644 --- a/tests/test_fp8_fp4.py +++ b/tests/test_fp8_fp4.py @@ -38,12 +38,11 @@ def test_gemm() -> None: a = a if major_a.is_k_major() else (a[0].T, a[1].T) b = b if major_b.is_k_major() else (b[0].T, b[1].T) assert a[0].is_contiguous() and b[0].is_contiguous() - a, a_sf = a - b, b_sf = b - getattr(deep_gemm, func_name)(a, a_sf, b, b_sf, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + getattr(deep_gemm, func_name)(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) diff = calc_diff(d, ref_d) assert diff < quant_config.max_diff(), (f'{m=}, {n=}, {k=}, {kernel_opt}, {major_opt=}, {accumulate=}, {out_dtype=}, ' f'{diff:.5f}, alias={test_alias}') + a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, kernel_type, use_ue8m0=use_ue8m0, quant_config=quant_config) t = bench_kineto(lambda: deep_gemm.fp8_fp4_gemm_nt(a, b, d, c=c, disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b), 'gemm_', suppress_kineto_output=True) @@ -217,7 +216,7 @@ def test_func(): print('Library path:') print(f' > {deep_gemm.__path__}\n') - # test_gemm() + test_gemm() test_m_grouped_gemm_contiguous() - # test_m_grouped_gemm_masked() - # test_k_grouped_gemm_contiguous() + test_m_grouped_gemm_masked() + test_k_grouped_gemm_contiguous() diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 5111edda23..2b007b6e60 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -56,6 +56,16 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): hidden, intermediate_hidden ) + # Cast weights into FP4 + def _cast_weights_to_fp4(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + num_groups, n, k = bf16_weights.shape + w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) + w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) + for i in range(num_groups): + w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) + w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) + return w, w_sf + # Create inputs # noinspection PyGlobalUndefined def create_inputs(): @@ -77,33 +87,12 @@ def create_inputs(): topk_idx.masked_fill_(rand_mask < args.masked_ratio, -1) topk_weights.masked_fill_(topk_idx < 0, 0) - # Check SF requirements - assert hidden % 128 == 0 - assert intermediate_hidden % 128 == 0 - assert l1_weights.shape[2] % 128 == 0 and l2_weights.shape[2] % 128 == 0 - - # Cast inputs to FP8 (or FP4 under DG_USE_FP4_ACTS) with per-32 UE8M0 SF. - # Stream A0.0b: when the flag is on, the symm buffer's `x` slot is sized - # for packed E2M1 (`hidden/2` bytes/token), so we must quantize at the - # source to match. - if os.environ.get('DG_USE_FP4_ACTS', '0') != '0': - x = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - else: - x = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - - # Cast grouped BF16 weights to FP4 with MN-major SF - # TODO: merge with `cast_fp8_fp4_with_major` - def cast_grouped_weights_to_fp4(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - num_groups, n, k = bf16_weights.shape - w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8) - w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float) - for i in range(num_groups): - w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32) - w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) - return w, w_sf - - l1_weights = cast_grouped_weights_to_fp4(l1_weights) - l2_weights = cast_grouped_weights_to_fp4(l2_weights) + # Cast inputs to FP8/FP4 with per-32 UE8M0 SF + assert hidden % 128 == 0 and intermediate_hidden % 128 == 0 + x = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + l1_weights = _cast_weights_to_fp4(l1_weights) + l2_weights = _cast_weights_to_fp4(l2_weights) + transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights) # Run fused mega MoE @@ -115,15 +104,13 @@ def run_fused(): buffer.topk_weights[:num_tokens].copy_(topk_weights) y = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device='cuda') - # noinspection PyTypeChecker - deep_gemm.fp8_fp4_mega_moe( - y, - transformed_l1_weights, transformed_l2_weights, - buffer, + kernel_kwargs = dict( + y=y, l1_weights=transformed_l1_weights, l2_weights=transformed_l2_weights, + sym_buffer=buffer, cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused, activation_clamp=args.activation_clamp, - fast_math=bool(args.fast_math) - ) + fast_math=bool(args.fast_math)) + deep_gemm.fp8_fp4_mega_moe(**kernel_kwargs) return y, cumulative_local_expert_recv_stats_fused dist_print('Config:', once_in_node=True) @@ -157,40 +144,44 @@ def run_fused(): num_topk=num_topk, use_fp8_dispatch=True, explicitly_destroy=True, allow_multiple_reduction=False, - gpu_timeout_secs=10, cpu_timeout_secs=30 + num_gpu_timeout_secs=10, num_cpu_timeout_secs=30 ) if is_legacy_loaded else None - def run_baseline(): - recv_x, _, recv_topk_weights, handle, _ = ep_buffer.dispatch( - x, topk_idx=topk_idx, topk_weights=topk_weights, - cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_baseline, - num_experts=num_experts, expert_alignment=alignment, - do_cpu_sync=False, do_handle_copy=False, - do_expand=True, use_tma_aligned_col_major_sf=True, - ) - n = recv_x[0].size(0) - l1_y = torch.empty((n, intermediate_hidden * 2), dtype=torch.bfloat16, device='cuda') - deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( - recv_x, l1_weights, l1_y, handle.psum_num_recv_tokens_per_expert, - use_psum_layout=True, recipe=(1, 1, 32)) - # noinspection PyCallingNonCallable - l1_y = tilelang_ops.swiglu_apply_weight_to_fp8( - x=l1_y, - topk_weights=recv_topk_weights, - avail_tokens=handle.psum_num_recv_tokens_per_expert[-1], - num_per_channels=32, - use_col_major_scales=True, - round_scale=True, - ue8m0_scale=True, - output_bf16=False, - clamp_value=args.activation_clamp, - fast_math=bool(args.fast_math) - ) - l2_y = torch.empty((n, hidden), dtype=torch.bfloat16, device='cuda') - deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( - l1_y, l2_weights, l2_y, handle.psum_num_recv_tokens_per_expert, - use_psum_layout=True, recipe=(1, 1, 32)) - return ep_buffer.combine(l2_y, handle=handle)[0], cumulative_local_expert_recv_stats_baseline + if is_legacy_loaded: + dispatch_kwargs = {'do_cpu_sync': False, 'do_handle_copy': False, + 'do_expand': True, 'use_tma_aligned_col_major_sf': True} + gemm_fn = deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous + gemm_kwargs = {'use_psum_layout': True, 'recipe': (1, 1, 32)} + activation_kwargs = {'round_scale': True, 'ue8m0_scale': True, 'output_bf16': False} + get_num_tokens = lambda recv_x: recv_x[0].size(0) + + def run_baseline(): + # Dispatch + recv_x, _, recv_topk_weights, handle, _ = ep_buffer.dispatch( + x, topk_idx=topk_idx, topk_weights=topk_weights, + cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_baseline, + num_experts=num_experts, expert_alignment=alignment, + **dispatch_kwargs) + num_recv_tokens = get_num_tokens(recv_x) + + # L1 GEMM + l1_y = torch.empty((num_recv_tokens, intermediate_hidden * 2), dtype=torch.bfloat16, device='cuda') + gemm_fn(recv_x, l1_weights, l1_y, handle.psum_num_recv_tokens_per_expert, **gemm_kwargs) + + # Activation + l1_y = tilelang_ops.swiglu_apply_weight_to_fp8( + x=l1_y, topk_weights=recv_topk_weights, + avail_tokens=handle.psum_num_recv_tokens_per_expert[-1], + num_per_channels=32, use_col_major_scales=True, + clamp_value=args.activation_clamp, fast_math=bool(args.fast_math), + **activation_kwargs) + + # L2 GEMM + l2_y = torch.empty((num_recv_tokens, hidden), dtype=torch.bfloat16, device='cuda') + gemm_fn(l1_y, l2_weights, l2_y, handle.psum_num_recv_tokens_per_expert, **gemm_kwargs) + + # Combine + return ep_buffer.combine(l2_y, handle=handle)[0], cumulative_local_expert_recv_stats_baseline # Check correctness (must be bitwise identical) num_correctness_tests = 1 if args.num_correctness_tests is None else args.num_correctness_tests @@ -224,15 +215,15 @@ def run_baseline(): safe_div = lambda a, b: float('nan') if b == 0 else a / b tflops = safe_div(2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_fused) - # HBM bytes: weights (FP4 packed = 0.5 bytes) + activations (FP8 = 1 byte) + output (BF16 = 2 bytes) - num_touched_experts = torch.unique(gathered_topk_idx.flatten()).numel() - 1 # NOTES minus 1 to exclude "-1" + # HBM bytes: weights + activations + output + num_touched_experts = torch.unique(gathered_topk_idx[gathered_topk_idx >= 0]).numel() num_hbm_bytes = ( - num_touched_experts * intermediate_hidden * 2 * hidden // 2 + # L1 weights (FP4) - num_touched_experts * hidden * intermediate_hidden // 2 + # L2 weights (FP4) - num_recv_tokens * hidden + # L1 acts read (FP8) - num_recv_tokens * intermediate_hidden + # L1 output write (FP8) - num_recv_tokens * intermediate_hidden + # L2 acts read (FP8) - num_recv_tokens * hidden * 2 # L2 output write (BF16) + num_touched_experts * intermediate_hidden * 2 * hidden * 0.5 # L1 weights + + num_touched_experts * hidden * intermediate_hidden * 0.5 # L2 weights + + num_recv_tokens * hidden # L1 acts read + + num_recv_tokens * intermediate_hidden # L1 output write + + num_recv_tokens * intermediate_hidden # L2 acts read + + num_recv_tokens * hidden * 2 # L2 output write (always BF16) ) hbm_gbs = safe_div(num_hbm_bytes / 1e9, t_fused) From 7020fd814ee1984f691ee8e2d2550ffe2ddf6601 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Wed, 10 Jun 2026 16:38:45 -0700 Subject: [PATCH 17/26] Change license in pyproject.toml to avoid build and publish failures with setuptools>=77 Co-authored-by: Brayden Zhong --- sgl_deep_gemm/pyproject.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sgl_deep_gemm/pyproject.toml b/sgl_deep_gemm/pyproject.toml index 94a0d22d4d..36d89a94a4 100644 --- a/sgl_deep_gemm/pyproject.toml +++ b/sgl_deep_gemm/pyproject.toml @@ -1,5 +1,5 @@ [build-system] -requires = ["setuptools>=61.0", "wheel"] +requires = ["setuptools>=77.0.0", "wheel"] build-backend = "setuptools.build_meta" [project] @@ -7,13 +7,13 @@ name = "sgl-deep-gemm" description = "SGLang fork of DeepGemm" readme = "README.md" requires-python = ">=3.10" -license = { file = "LICENSE" } +license = "Apache-2.0" +license-files = ["LICENSE"] authors = [ {name = "SGLang Kernel Team", email="sglang@lmsys.org"}, ] classifiers = [ "Programming Language :: Python :: 3", - "License :: OSI Approved :: Apache Software License", "Environment :: GPU :: NVIDIA CUDA" ] dynamic = ["version"] From c872baf26eab453cc2dd426e2e0e4b77c54703cd Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Wed, 10 Jun 2026 19:03:35 -0700 Subject: [PATCH 18/26] Move run tests script for DeepGEMM (#42) --- sgl_deep_gemm/run_tests.sh | 182 +++++++++++++++++++++++++++++++++++++ 1 file changed, 182 insertions(+) create mode 100755 sgl_deep_gemm/run_tests.sh diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh new file mode 100755 index 0000000000..20b5394a6b --- /dev/null +++ b/sgl_deep_gemm/run_tests.sh @@ -0,0 +1,182 @@ +#!/usr/bin/env bash +# Pre-release gate for the installed sgl-deep-gemm wheel (sglang's release-whl-deepgemm.yml). +# Runs from sgl_deep_gemm/tests/ so `import deep_gemm` hits the wheel's prebuilt _C.so, not +# the source tree (which JIT-rebuilds _C); the guard below aborts if that resolution is wrong. +# +# Usage: run_tests.sh [DEEPGEMM_SRC] [--max-procs N] [--skip-sanitizer] [--skip-mega-moe] +# DEEPGEMM_SRC defaults to this script's own repo root. +set -uo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +DEEPGEMM_SRC="$(cd "${SCRIPT_DIR}/.." && pwd)" +if [ $# -gt 0 ] && [ "${1#--}" = "$1" ]; then + DEEPGEMM_SRC="$(cd "$1" && pwd)"; shift +fi +MAX_PROCS="" +SKIP_SANITIZER=0 +SKIP_MEGA_MOE=0 +while [ $# -gt 0 ]; do + case "$1" in + --max-procs) MAX_PROCS="$2"; shift 2 ;; + --skip-sanitizer) SKIP_SANITIZER=1; shift ;; + --skip-mega-moe) SKIP_MEGA_MOE=1; shift ;; + *) echo "Unknown option: $1" >&2; exit 1 ;; + esac +done + +TESTS_DIR="${DEEPGEMM_SRC}/sgl_deep_gemm/tests" +PYTHON="${PYTHON:-python3}" + +if [ ! -d "${TESTS_DIR}" ]; then + echo "No sgl_deep_gemm/tests/ directory under ${DEEPGEMM_SRC}" >&2 + exit 1 +fi + +NUM_GPUS=$(nvidia-smi -L 2>/dev/null | wc -l) +if [ "${NUM_GPUS}" -eq 0 ]; then + echo "No GPUs visible to nvidia-smi — DeepGEMM tests require a GPU." >&2 + exit 1 +fi +# arch major: 9 == Hopper (SM90), 10 == Blackwell (SM100). +COMPUTE_CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1) +ARCH_MAJOR=${COMPUTE_CAP%%.*} + +NPROC=${NUM_GPUS} +if [ -n "${MAX_PROCS}" ] && [ "${MAX_PROCS}" -lt "${NPROC}" ]; then + NPROC=${MAX_PROCS} +fi + +DG_FILE=$(cd "${TESTS_DIR}" && "${PYTHON}" -c "import deep_gemm, sys; sys.stdout.write(deep_gemm.__file__)" 2>/dev/null) +if [ -z "${DG_FILE}" ]; then + echo "Failed to import deep_gemm — is the wheel installed?" >&2 + exit 1 +fi +case "${DG_FILE}" in + "${DEEPGEMM_SRC}"/*) + echo "ERROR: 'import deep_gemm' resolved to the source tree (${DG_FILE})," >&2 + echo " not the installed wheel. Aborting." >&2 + exit 1 ;; +esac + +echo "==============================================================" +echo " DeepGEMM wheel test run" +echo " deep_gemm: ${DG_FILE}" +echo " GPUs: ${NUM_GPUS} (compute_cap ${COMPUTE_CAP}, arch major ${ARCH_MAJOR})" +echo " processes: ${NPROC}" +echo " skip-mega: ${SKIP_MEGA_MOE} skip-sanitizer: ${SKIP_SANITIZER}" +echo "==============================================================" + +PASSED=() +FAILED=() +SKIPPED=() + +run_test() { + local name="$1"; shift + echo "" + echo "----- RUN ${name} $* -----" + local log; log=$(mktemp) + (cd "${TESTS_DIR}" && "${PYTHON}" "${name}" "$@") 2>&1 | tee "${log}" + local rc=${PIPESTATUS[0]} + # Fork-based multiprocessing tests can crash in child processes while the + # launcher still exits 0; treat an unhandled traceback as a failure too. + if [ "${rc}" -eq 0 ] && grep -q "Traceback (most recent call last)" "${log}"; then + echo "----- FAIL ${name} (child process crashed; launcher exited 0) -----" + FAILED+=("${name}") + elif [ "${rc}" -eq 0 ]; then + echo "----- PASS ${name} -----" + PASSED+=("${name}") + else + echo "----- FAIL ${name} (exit ${rc}) -----" + FAILED+=("${name}") + fi + rm -f "${log}" +} + +skip_test() { + echo "" + echo "----- SKIP $1 ($2) -----" + SKIPPED+=("$1 ($2)") +} + +# test_legacy.py is intentionally excluded: the deep_gemm.legacy kernels are +# deprecated and not exposed by the wheel. +SINGLE_GPU_TESTS=( + test_bf16.py + test_einsum.py + test_fp8_fp4.py + test_hyperconnection.py + test_layout.py + test_attention.py +) +for t in "${SINGLE_GPU_TESTS[@]}"; do + if [ -f "${TESTS_DIR}/${t}" ]; then + run_test "${t}" + else + skip_test "${t}" "not present in this branch" + fi +done + +# test_lazy_init.py is intentionally excluded: `import tvm_ffi` eagerly creates a +# CUDA context, so `import deep_gemm` trips torch's bad-fork guard. Tracked +# upstream in apache-tvm-ffi; re-enable once import no longer initializes CUDA. +if [ -f "${TESTS_DIR}/test_lazy_init.py" ]; then + skip_test test_lazy_init.py "tvm_ffi eager CUDA init (upstream)" +fi + +# mega_moe family uses SM100 fp4 + symmetric-memory kernels (SM100-only). +# test_mega_moe.py additionally needs deep_ep (with ElasticBuffer); the l1 and +# pre_dispatch tests use deep_gemm's own symmetric buffer. +MEGA_MOE_ALL=( + test_mega_moe.py + test_mega_moe_l1_fp4_accuracy.py + test_mega_moe_l1_sentinel.py + test_mega_moe_pre_dispatch.py +) +MEGA_MOE_L1=( + test_mega_moe_l1_fp4_accuracy.py + test_mega_moe_l1_sentinel.py +) +if [ "${SKIP_MEGA_MOE}" -eq 1 ]; then + for t in "${MEGA_MOE_ALL[@]}"; do + [ -f "${TESTS_DIR}/${t}" ] && skip_test "${t}" "--skip-mega-moe" + done +elif [ "${ARCH_MAJOR}" -ge 10 ]; then + if [ -f "${TESTS_DIR}/test_mega_moe.py" ]; then + if (cd "${TESTS_DIR}" && "${PYTHON}" -c "import deep_ep; assert hasattr(deep_ep, 'ElasticBuffer')") >/dev/null 2>&1; then + run_test test_mega_moe.py --num-processes "${NPROC}" + else + skip_test test_mega_moe.py "deep_ep with ElasticBuffer not installed" + fi + fi + # l1 tests are quarantined: they exercise a manual buffer-packing path that + # diverges from sglang's pre_dispatch flow, and hit kernel-level fp4 failures + # (TMA stride at >=8 ranks, rel-RMSE). sglang's real path is covered by + # test_mega_moe_pre_dispatch + test_mega_moe. Confirm on B200 before re-enabling. + for t in "${MEGA_MOE_L1[@]}"; do + [ -f "${TESTS_DIR}/${t}" ] && skip_test "${t}" "fp4 kernel failures, see comment (confirm on B200)" + done + [ -f "${TESTS_DIR}/test_mega_moe_pre_dispatch.py" ] && run_test test_mega_moe_pre_dispatch.py +else + for t in "${MEGA_MOE_ALL[@]}"; do + [ -f "${TESTS_DIR}/${t}" ] && skip_test "${t}" "SM100-only, arch major ${ARCH_MAJOR}" + done +fi + +# test_sanitizer.py is intentionally excluded: compute-sanitizer memcheck/synccheck +# are clean, but its DG_JIT_PTXAS_CHECK trips on a register spill ("Local memory +# used") in fp8_fp4_mqa_logits — a perf/codegen finding, not a memory-safety bug. +# Re-enable once that kernel's register pressure is addressed or the check is +# scoped to allow it. +if [ -f "${TESTS_DIR}/test_sanitizer.py" ]; then + skip_test test_sanitizer.py "fp8_fp4_mqa_logits register spill (known)" +fi + +echo "" +echo "==============================================================" +echo " Summary: ${#PASSED[@]} passed, ${#FAILED[@]} failed, ${#SKIPPED[@]} skipped" +[ ${#PASSED[@]} -gt 0 ] && printf ' PASS %s\n' "${PASSED[@]}" +[ ${#SKIPPED[@]} -gt 0 ] && printf ' SKIP %s\n' "${SKIPPED[@]}" +[ ${#FAILED[@]} -gt 0 ] && printf ' FAIL %s\n' "${FAILED[@]}" +echo "==============================================================" + +[ ${#FAILED[@]} -eq 0 ] From 2176dff556d6fb8f0184516dc224a3b31a094b98 Mon Sep 17 00:00:00 2001 From: Yuxuan Zhang Date: Sat, 13 Jun 2026 05:17:56 +0800 Subject: [PATCH 19/26] Support num_heads == 16 in MQA logits (#43) Add host-side assertions and fix SM100 FP8/FP4 (paged) MQA logits kernels for num_heads == 16. --- csrc/apis/attention.hpp | 8 ++++---- .../impls/smxx_fp8_fp4_paged_mqa_logits.hpp | 6 ++++-- .../deep_gemm/impls/sm100_fp4_mqa_logits.cuh | 9 +++++---- .../deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh | 14 +++++++++----- .../deep_gemm/impls/sm100_fp8_mqa_logits.cuh | 9 +++++---- .../deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh | 14 +++++++++----- 6 files changed, 36 insertions(+), 24 deletions(-) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index cecace56ce..49f5c93738 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -90,7 +90,7 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple(q_fp); head_dim *= 2; - DG_HOST_ASSERT(num_heads == 32 or num_heads == 64); + DG_HOST_ASSERT(num_heads == 16 or num_heads == 32 or num_heads == 64); DG_HOST_ASSERT(head_dim == 128); DG_HOST_ASSERT(q_fp.is_contiguous()); DG_HOST_ASSERT(q_fp.scalar_type() == kPackedFP4); @@ -117,7 +117,7 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple(q_fp); - DG_HOST_ASSERT(num_heads == 32 or num_heads == 64); + DG_HOST_ASSERT(num_heads == 16 or num_heads == 32 or num_heads == 64); DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128); DG_HOST_ASSERT(q_fp.is_contiguous()); DG_HOST_ASSERT(q_fp.scalar_type() == torch::kFloat8_e4m3fn); @@ -247,7 +247,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple(q_fp); head_dim *= 2; DG_HOST_ASSERT(next_n >= 1); - DG_HOST_ASSERT(num_heads == 32 or num_heads == 64); + DG_HOST_ASSERT(num_heads == 16 or num_heads == 32 or num_heads == 64); DG_HOST_ASSERT(head_dim == 128); DG_HOST_ASSERT(q_fp.is_contiguous()); DG_HOST_ASSERT(q_fp.scalar_type() == kPackedFP4); @@ -285,7 +285,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple(q_fp); DG_HOST_ASSERT(next_n >= 1); - DG_HOST_ASSERT(num_heads == 32 or num_heads == 64); + DG_HOST_ASSERT(num_heads == 16 or num_heads == 32 or num_heads == 64); DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128); DG_HOST_ASSERT(q_fp.is_contiguous()); DG_HOST_ASSERT(q_fp.scalar_type() == torch::kFloat8_e4m3fn); diff --git a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp index 2a3288ee0c..0033baa086 100644 --- a/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp @@ -235,7 +235,8 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q, const int smem_q_size_per_stage = next_n_atom * num_heads * head_dim * static_cast(q.element_size()); const int smem_kv_size_per_stage = split_kv * head_dim * static_cast(kv_cache.element_size()); const int smem_kv_scale_size_per_stage = split_kv * static_cast(kv_cache_scales.element_size()); - const int smem_weight_size_per_stage = next_n_atom * num_heads * static_cast(weights.element_size()); + // The weight stage stride is padded to 128B in the kernel for TMA alignment + const int smem_weight_size_per_stage = align(next_n_atom * num_heads * static_cast(weights.element_size()), 128); const int smem_barriers = (num_q_stages + num_kv_stages) * 2 * 8; const int smem_umma_barriers = num_math_warp_groups * 2 * 8; @@ -415,7 +416,8 @@ static void sm100_fp4_paged_mqa_logits(const torch::Tensor& q, const int smem_sf_q_size_per_stage = align(next_n_atom * num_heads, 128) * sizeof(int); const int smem_kv_size_per_stage = split_kv * head_dim / 2; const int smem_sf_kv_size_per_stage = align(split_kv, 128) * sizeof(int); - const int smem_weight_size_per_stage = next_n_atom * num_heads * sizeof(float); + // The weight stage stride is padded to 128B in the kernel for TMA alignment + const int smem_weight_size_per_stage = align(next_n_atom * num_heads * static_cast(sizeof(float)), 128); const int smem_barriers = (num_q_stages + num_kv_stages + num_tmem_stages) * 2 * 8; const int smem_tmem_ptr = 4; diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp4_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp4_mqa_logits.cuh index b8a99fd042..54d638750b 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp4_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp4_mqa_logits.cuh @@ -349,10 +349,11 @@ void sm100_fp4_mqa_logits(const uint32_t seq_len, const uint32_t seq_len_kv, // Helper lambda for loading tensor memory auto tmem_load = [](auto num_elems_c, const uint32_t& tmem_addr, float* accum) { constexpr uint32_t N = decltype(num_elems_c)::value; - DG_STATIC_ASSERT(N == 32 or N == 64, "Unsupported TMEM load size"); - using Loader = cute::conditional_t; + DG_STATIC_ASSERT(N == 8 or N == 16 or N == 32 or N == 64, "Unsupported TMEM load size"); + using Loader = cute::conditional_t>>; [&](cute::index_sequence) { Loader::copy(tmem_addr, reinterpret_cast(accum)[Is]...); }(cute::make_index_sequence{}); diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh index d9add53425..1183a97d57 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh @@ -79,6 +79,9 @@ void sm100_fp4_paged_mqa_logits(const uint32_t batch_size, static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = SPLIT_KV * (kHeadDim / 2); static constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = kNumSFKV * sizeof(int); static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = kNextNAtom * kNumHeads * sizeof(float); + // TMA requires 128B-aligned shared memory addresses, so pad the per-stage stride + // (e.g. `kNextNAtom == 1` with 16 heads gives only 64B per stage) + static constexpr uint32_t SMEM_WEIGHT_STRIDE_PER_STAGE = SMEM_WEIGHT_SIZE_PER_STAGE < 128 ? 128 : SMEM_WEIGHT_SIZE_PER_STAGE; // Align to swizzling alignment bytes extern __shared__ __align__(kSwizzleAlignment) uint8_t smem_buffer[]; @@ -101,7 +104,7 @@ void sm100_fp4_paged_mqa_logits(const uint32_t batch_size, }); auto smem_weights = utils::PatternVisitor([&](const uint32_t& i) { return reinterpret_cast(smem_sf_ptr + SMEM_SF_Q_SIZE_PER_STAGE * kNumQStages + SMEM_SF_KV_SIZE_PER_STAGE * kNumKVStages - + SMEM_WEIGHT_SIZE_PER_STAGE * i); + + SMEM_WEIGHT_STRIDE_PER_STAGE * i); }); // Barriers and TMEM pointer on shared memory @@ -382,10 +385,11 @@ void sm100_fp4_paged_mqa_logits(const uint32_t batch_size, // Helper lambda for loading tensor memory auto tmem_load = [](auto num_elems_c, const uint32_t& tmem_addr, float* accum) { constexpr int N = decltype(num_elems_c)::value; - DG_STATIC_ASSERT(N == 32 or N == 64, "Unsupported TMEM load size"); - using Loader = cute::conditional_t; + DG_STATIC_ASSERT(N == 8 or N == 16 or N == 32 or N == 64, "Unsupported TMEM load size"); + using Loader = cute::conditional_t>>; [&](cute::index_sequence) { Loader::copy(tmem_addr, reinterpret_cast(accum)[Is]...); }(cute::make_index_sequence{}); diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh index e6744f59ac..5e3699c85f 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_mqa_logits.cuh @@ -293,10 +293,11 @@ void sm100_fp8_mqa_logits(const uint32_t seq_len, const uint32_t seq_len_kv, // Helper lambda for loading tensor memory auto tmem_load = [](auto num_elems_c, const uint32_t& tmem_addr, float* accum) { constexpr int N = decltype(num_elems_c)::value; - DG_STATIC_ASSERT(N == 32 or N == 64, "Unsupported TMEM load size"); - using Loader = cute::conditional_t; + DG_STATIC_ASSERT(N == 8 or N == 16 or N == 32 or N == 64, "Unsupported TMEM load size"); + using Loader = cute::conditional_t>>; [&](cute::index_sequence) { Loader::copy(tmem_addr, reinterpret_cast(accum)[Is]...); }(cute::make_index_sequence{}); diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh index 9a5bddbf37..e63d56c9fb 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh @@ -65,6 +65,9 @@ void sm100_fp8_paged_mqa_logits(const uint32_t batch_size, static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = SPLIT_KV * kHeadDim * sizeof(__nv_fp8_e4m3); static constexpr uint32_t SMEM_KV_SCALE_SIZE_PER_STAGE = SPLIT_KV * sizeof(float); static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = kNextNAtom * kNumHeads * sizeof(float); + // TMA requires 128B-aligned shared memory addresses, so pad the per-stage stride + // (e.g. `kNextNAtom == 1` with 16 heads gives only 64B per stage) + static constexpr uint32_t SMEM_WEIGHT_STRIDE_PER_STAGE = SMEM_WEIGHT_SIZE_PER_STAGE < 128 ? 128 : SMEM_WEIGHT_SIZE_PER_STAGE; // Align to swizzling alignment bytes extern __shared__ __align__(kSwizzleAlignment) uint8_t smem_buffer[]; @@ -83,7 +86,7 @@ void sm100_fp8_paged_mqa_logits(const uint32_t batch_size, return reinterpret_cast(smem_buffer + smem_offset + SMEM_KV_SCALE_SIZE_PER_STAGE * i); }); auto smem_weights = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + smem_offset + SMEM_KV_SCALE_SIZE_PER_STAGE * kNumKVStages + SMEM_WEIGHT_SIZE_PER_STAGE * i); + return reinterpret_cast(smem_buffer + smem_offset + SMEM_KV_SCALE_SIZE_PER_STAGE * kNumKVStages + SMEM_WEIGHT_STRIDE_PER_STAGE * i); }); // Barriers and TMEM pointer on shared memory @@ -307,10 +310,11 @@ void sm100_fp8_paged_mqa_logits(const uint32_t batch_size, // Helper lambda for loading tensor memory auto tmem_load = [](auto num_elems_c, const uint32_t& tmem_addr, float* accum) { constexpr int N = decltype(num_elems_c)::value; - DG_STATIC_ASSERT(N == 32 or N == 64, "Unsupported TMEM load size"); - using Loader = cute::conditional_t; + DG_STATIC_ASSERT(N == 8 or N == 16 or N == 32 or N == 64, "Unsupported TMEM load size"); + using Loader = cute::conditional_t>>; [&](cute::index_sequence) { Loader::copy(tmem_addr, reinterpret_cast(accum)[Is]...); }(cute::make_index_sequence{}); From 35d4d8c34927f8ffc4f4e8957624ae8c87eb7c6d Mon Sep 17 00:00:00 2001 From: Ding Yin Date: Tue, 16 Jun 2026 07:05:46 +0800 Subject: [PATCH 20/26] Sm90 mega moe on sgl dev (#36) Co-authored-by: yinding --- csrc/apis/sm90_mega.hpp | 204 ++ csrc/jit_kernels/heuristics/sm90_mega_moe.hpp | 191 ++ csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp | 305 +++ csrc/tvm_ffi_api.cpp | 60 + deep_gemm/include/deep_gemm/comm/barrier.cuh | 5 + .../deep_gemm/impls/sm90_fp8_mega_moe.cuh | 1935 ++++++++++++++++ sgl_deep_gemm/__init__.py | 139 +- tests/test_mega_moe_hopper.py | 1946 +++++++++++++++++ 8 files changed, 4783 insertions(+), 2 deletions(-) create mode 100644 csrc/apis/sm90_mega.hpp create mode 100644 csrc/jit_kernels/heuristics/sm90_mega_moe.hpp create mode 100644 csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp create mode 100644 deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh create mode 100644 tests/test_mega_moe_hopper.py diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp new file mode 100644 index 0000000000..9bfe73c91a --- /dev/null +++ b/csrc/apis/sm90_mega.hpp @@ -0,0 +1,204 @@ +#pragma once + +#include + +#include "mega.hpp" +#include "../jit_kernels/impls/sm90_fp8_mega_moe.hpp" + +namespace deep_gemm::mega { + +static int get_token_alignment_for_sm90_mega_moe() { + return layout::kLCMCandidateBlockM; +} + +static std::tuple(const torch::Tensor&)>> +get_symm_buffer_size_for_sm90_mega_moe( + const int& num_ranks, const int& num_experts, + const int& num_max_tokens_per_rank, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const bool& use_fp8_dispatch, const std::string& activation) { + DG_HOST_ASSERT(num_experts % num_ranks == 0); + DG_HOST_ASSERT(use_fp8_dispatch); + DG_HOST_ASSERT(activation == "swiglu"); + + const auto workspace = layout::Workspace(nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk); + + const auto fp8_token_layout = layout::Data(hidden); + const auto bf16_token_layout = layout::Data(hidden * 2); + const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden); + const auto fp8_sf_layout = layout::Data(hidden / 32); + const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 16); + const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false); + const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false); + const auto l1_topk_weights_layout = layout::Data(sizeof(float), false); + + const auto input_token_buffer = layout::Buffer( + fp8_token_layout, 1, num_max_tokens_per_rank, + workspace.get_end_ptr()); + const auto input_sf_buffer = layout::Buffer( + fp8_sf_layout, 1, num_max_tokens_per_rank, + input_token_buffer.get_end_ptr()); + const auto input_topk_idx_buffer = layout::Buffer( + input_topk_idx_layout, 1, num_max_tokens_per_rank, + input_sf_buffer.get_end_ptr()); + const auto input_topk_weights_buffer = layout::Buffer( + input_topk_weights_layout, 1, num_max_tokens_per_rank, + input_topk_idx_buffer.get_end_ptr()); + + const auto num_max_pool_tokens = static_cast(workspace.num_max_pool_tokens); + int num_max_padded_sf_pool_tokens = 0; + for (int block_m: layout::kCandidateBlockM) { + num_max_padded_sf_pool_tokens = std::max( + num_max_padded_sf_pool_tokens, + layout::get_num_padded_sf_pool_tokens(num_max_pool_tokens, block_m) + ); + } + + const auto l1_token_buffer = layout::Buffer( + fp8_token_layout, 1, num_max_pool_tokens, + input_topk_weights_buffer.get_end_ptr()); + const auto l1_sf_buffer = layout::Buffer( + fp8_sf_layout, 1, num_max_padded_sf_pool_tokens, + l1_token_buffer.get_end_ptr()); + const auto l1_topk_weights_buffer = layout::Buffer( + l1_topk_weights_layout, 1, num_max_pool_tokens, + l1_sf_buffer.get_end_ptr()); + + const auto l2_token_buffer = layout::Buffer( + fp8_intermediate_token_layout, 1, num_max_pool_tokens, + l1_topk_weights_buffer.get_end_ptr()); + const auto l2_sf_buffer = layout::Buffer( + fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens, + l2_token_buffer.get_end_ptr()); + + const auto combine_token_buffer = layout::Buffer( + bf16_token_layout, num_topk, num_max_tokens_per_rank, + l2_sf_buffer.get_end_ptr()); + + DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); + + auto slice_input_buffers = [=](const torch::Tensor& buffer) { + auto x = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_token_buffer.base)), + {num_max_tokens_per_rank, hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto x_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_sf_buffer.base)), + {num_max_tokens_per_rank, hidden / 128}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto topk_idx = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_topk_idx_buffer.base)), + {num_max_tokens_per_rank, num_topk}, + torch::TensorOptions().dtype(torch::kInt64).device(buffer.device())); + auto topk_weights = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(input_topk_weights_buffer.base)), + {num_max_tokens_per_rank, num_topk}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto l1_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_token_buffer.base)), + {num_max_pool_tokens, hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto l1_acts_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l1_sf_buffer.base)), + {num_max_padded_sf_pool_tokens, hidden / 128}, + {1, num_max_padded_sf_pool_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + auto l2_acts = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_token_buffer.base)), + {num_max_pool_tokens, intermediate_hidden}, + torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device())); + auto l2_acts_sf = torch::from_blob( + math::advance_ptr(buffer.data_ptr(), reinterpret_cast(l2_sf_buffer.base)), + {num_max_padded_sf_pool_tokens, intermediate_hidden / 64}, + {1, num_max_padded_sf_pool_tokens}, + torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device())); + return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf); + }; + return {reinterpret_cast(combine_token_buffer.get_end_ptr()), slice_input_buffers}; +} + +static void fp8_mega_moe( + const torch::Tensor& y, + const std::tuple& l1_weights_tuple, + const std::tuple& l2_weights_tuple, + const std::optional& cumulative_local_expert_recv_stats, + const torch::Tensor& sym_buffer, + const std::vector& sym_buffer_ptrs, const int& rank_idx, + const int& num_max_tokens_per_rank, + const int& num_experts, const int& num_topk, + const std::tuple& recipe, + const std::string& activation, + const std::optional& activation_clamp_opt, + const bool& fast_math +) { + const auto [l1_weights, l1_weights_sf] = l1_weights_tuple; + const auto [l2_weights, l2_weights_sf] = l2_weights_tuple; + + const auto arch_major = device_runtime->get_arch_major(); + DG_HOST_ASSERT(arch_major == 9); + + const auto num_tokens = static_cast(y.size(0)); + const auto [rm, rn, rk] = recipe; + DG_HOST_ASSERT(rm == 128 and rn == 128 and rk == 128); + DG_HOST_ASSERT(activation == "swiglu"); + + const auto activation_clamp = + activation_clamp_opt.value_or(std::numeric_limits::infinity()); + DG_HOST_ASSERT(activation_clamp >= 0); + + DG_HOST_ASSERT(get_major_type_ab(l1_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(get_major_type_ab(l2_weights) == cute::UMMA::Major::K); + DG_HOST_ASSERT(l1_weights.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(l2_weights.scalar_type() == torch::kFloat8_e4m3fn); + const auto [num_experts_per_rank, intermediate_hidden_2, hidden] = get_shape<3>(l1_weights); + const auto [num_experts_per_rank_, hidden_, intermediate_hidden] = get_shape<3>(l2_weights); + DG_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank); + DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_); + DG_HOST_ASSERT(hidden == hidden_); + DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden); + DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous()); + DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0); + DG_HOST_ASSERT(intermediate_hidden / 64 <= 64); + + constexpr int kGranMN = 128, kGranK = 128; + check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK, + num_experts_per_rank, false, true, torch::kFloat); + check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK, + num_experts_per_rank, false, true, torch::kFloat); + + if (cumulative_local_expert_recv_stats.has_value()) { + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank); + DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous()); + } + + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts_ = num_experts_per_rank * num_ranks; + const auto [num_required_bytes, slice] = get_symm_buffer_size_for_sm90_mega_moe( + num_ranks, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + true, activation); + DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); + DG_HOST_ASSERT(num_experts == num_experts_); + + const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer); + + sm90_fp8_mega_moe(y, + l1_acts, l1_acts_sf, + l2_acts, l2_acts_sf, + l1_weights, l2_weights, + l1_weights_sf, l2_weights_sf, + cumulative_local_expert_recv_stats, + sym_buffer_ptrs, + rank_idx, num_max_tokens_per_rank, + num_experts_per_rank, + num_tokens, num_topk, + hidden, intermediate_hidden, + activation_clamp, fast_math); + + if (get_env("DG_COMM_KERNEL_DEBUG")) + sym_buffer.zero_(); +} + +} // namespace deep_gemm::mega diff --git a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp new file mode 100644 index 0000000000..b62f85b268 --- /dev/null +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -0,0 +1,191 @@ +#pragma once + +#include "mega_moe.hpp" + +namespace deep_gemm { + +// ============================================================================ +// SM90 (Hopper) MegaMoE configuration +// ---------------------------------------------------------------------------- +// SM90 differs from SM100 in: +// - No tensor memory (TMEM): WGMMA accumulators live in registers. +// - No FP4: weights are FP8 e4m3 with per-128 channel float scales. +// - No 2-CTA cluster MMA: TMA multicast cluster=2 may still be used. +// - Activation SF is float, not UE8M0 int: L1 input uses per-128 K and the +// fused L1 epilogue writes L2 activation SF at per-64 K granularity. +// The kernel implementation is in `deep_gemm/impls/sm90_fp8_mega_moe.cuh`. +// ============================================================================ + +struct MegaMoESM90Config { + int block_m, block_n, block_k; + int cluster_size; + int num_max_pool_tokens; + int num_padded_sf_pool_tokens; + int swizzle_acts_mode, swizzle_weights_mode; + int num_experts_per_wave; + int num_stages, smem_size; + int num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads; + + friend std::ostream& operator << (std::ostream& os, const MegaMoESM90Config& config) { + os << "MegaMoESM90Config(" + << "block_m=" << config.block_m << ", block_n=" << config.block_n << ", block_k=" << config.block_k + << ", cluster_size=" << config.cluster_size + << ", num_max_pool_tokens=" << config.num_max_pool_tokens + << ", num_padded_sf_pool_tokens=" << config.num_padded_sf_pool_tokens + << ", swizzle_acts_mode=" << config.swizzle_acts_mode << ", swizzle_weights_mode=" << config.swizzle_weights_mode + << ", num_experts_per_wave=" << config.num_experts_per_wave + << ", num_stages=" << config.num_stages << ", smem_size=" << config.smem_size + << ", num_dispatch_threads=" << config.num_dispatch_threads + << ", num_non_epilogue_threads=" << config.num_non_epilogue_threads + << ", num_epilogue_threads=" << config.num_epilogue_threads << ")"; + return os; + } +}; + +static std::tuple get_block_config_for_mega_moe_sm90( + const int& num_ranks, const int& num_experts, + const int& num_max_tokens_per_rank, const int& num_topk, + const int& num_tokens) { + const float expected_tokens_per_expert = + static_cast(num_tokens) * num_ranks * num_topk / num_experts; + const bool auto_split_mn = expected_tokens_per_expert >= 64.0f; + if (auto_split_mn) + return {128, 512}; + + const int block_m = 64; + const int num_epilogue_warpgroups = 2; + + DG_HOST_ASSERT(std::any_of( + layout::kCandidateBlockM, layout::kCandidateBlockM + layout::kNumCandidateBlockMs, + [=](const auto& candidate) { return candidate == block_m; }) + ); + return {block_m, num_epilogue_warpgroups * 128}; +} + +static int get_num_experts_per_wave_for_mega_moe_sm90( + const int& num_experts_per_rank, const int& num_tokens, const int& num_topk, + const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms) { + const float expected_tokens_per_expert = + static_cast(num_tokens) * num_topk / num_experts_per_rank; + if (expected_tokens_per_expert < 1.0f or expected_tokens_per_expert > 4.0f) + return num_experts_per_rank; + + if (block_m == 64 and intermediate_hidden >= 3072) { + const int num_n_blocks_per_expert = (2 * intermediate_hidden) / block_n; + const int single_wave_blocks = + num_experts_per_rank * num_n_blocks_per_expert; + if (single_wave_blocks >= 4 * num_sms) + return num_experts_per_rank; + } + return get_num_experts_per_wave_for_mega_moe( + num_experts_per_rank, num_tokens, num_topk, + intermediate_hidden, block_m, block_n, num_sms); +} + +static std::pair get_pipeline_config_for_mega_moe_sm90( + const int& smem_capacity, + const int& num_experts, const int& hidden, + const int& block_m, const int& block_n, const int& block_k, + const int& num_dispatch_warps, const int& num_epilogue_warps) { + constexpr int kSmemAlignment = 1024; + + const int smem_expert_count_size = align( + num_experts * static_cast(sizeof(uint32_t)), kSmemAlignment); + const int smem_send_buffers_size = align( + static_cast(layout::Buffer(layout::Data(hidden), num_dispatch_warps, 1).get_num_bytes()), + kSmemAlignment); + const int smem_dispatch_size = smem_expert_count_size + smem_send_buffers_size; + + const int smem_cd_l1 = block_m * (block_n / 2); + const int smem_cd_l2 = block_m * block_n * static_cast(sizeof(nv_bfloat16)); + const int smem_cd = align(std::max(smem_cd_l1, smem_cd_l2), kSmemAlignment); + + const int smem_sfa_per_stage = align(2 * block_m * static_cast(sizeof(float)), 128); + const int smem_sfb_per_stage = 0; + const int smem_per_stage = block_m * block_k + block_n * block_k + + smem_sfa_per_stage + smem_sfb_per_stage; + + const int smem_barriers_fixed = (num_dispatch_warps + 2 * num_epilogue_warps) * 8; + const int smem_barriers_per_stage = 2 * 8; + const int smem_fixed = smem_dispatch_size + smem_cd + smem_barriers_fixed; + + const int num_stages = (smem_capacity - smem_fixed) / + (smem_per_stage + smem_barriers_per_stage); + DG_HOST_ASSERT(num_stages >= 2); + const int smem_size = smem_fixed + num_stages * (smem_per_stage + smem_barriers_per_stage); + DG_HOST_ASSERT(smem_size <= smem_capacity); + return {num_stages, smem_size}; +} + +static MegaMoESM90Config get_mega_moe_config_sm90( + const int& num_ranks, const int& num_experts, const int& num_experts_per_rank, + const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const int& num_padded_sf_pool_tokens) { + const auto [block_m, num_epilogue_threads] = get_block_config_for_mega_moe_sm90( + num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens); + const float expected_tokens_per_expert = + static_cast(num_tokens) * num_ranks * num_topk / num_experts; + const bool auto_split_mn = expected_tokens_per_expert >= 64.0f; + const bool decode_split_n_path = + block_m == 64 and num_epilogue_threads == 256; + const bool decode_use_block_n_256 = + decode_split_n_path and intermediate_hidden >= 3072 and + expected_tokens_per_expert >= 0.25f and + (2 * intermediate_hidden) % 256 == 0; + const int block_n = auto_split_mn ? 256 + : (decode_use_block_n_256 ? 256 : 128); + const int block_k = 128; + const int cluster_size = 1; + const int num_max_pool_tokens = layout::get_num_max_pool_tokens( + num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank); + const int swizzle_acts_mode = 128; + const int swizzle_weights_mode = 128; + + const int num_sms = device_runtime->get_num_sms(); + const int num_experts_per_wave = get_num_experts_per_wave_for_mega_moe_sm90( + num_experts_per_rank, num_tokens, num_topk, + intermediate_hidden, block_m, block_n, num_sms); + + const bool reduce_decode_threads = num_epilogue_threads == 128; + const bool decode_split_n = + block_m == 64 and num_epilogue_threads == 256; + const bool shrink_non_epilogue = reduce_decode_threads or decode_split_n; + const int num_dispatch_threads = + (num_epilogue_threads == 512 or shrink_non_epilogue) ? 64 : 128; + const bool split_sfa_loader_warp = false; + const int num_non_epilogue_threads = + split_sfa_loader_warp ? 128 : + ((num_epilogue_threads == 512 or shrink_non_epilogue) ? 64 : 128); + DG_HOST_ASSERT((num_dispatch_threads + num_non_epilogue_threads) % 128 == 0); + + const auto [num_stages, smem_size] = get_pipeline_config_for_mega_moe_sm90( + SM90ArchSpec::smem_capacity, + num_experts, hidden, + block_m, block_n, block_k, + num_dispatch_threads / 32, num_epilogue_threads / 32); + + const auto config = MegaMoESM90Config { + block_m, block_n, block_k, + cluster_size, + num_max_pool_tokens, num_padded_sf_pool_tokens, + swizzle_acts_mode, swizzle_weights_mode, + num_experts_per_wave, + num_stages, smem_size, + num_dispatch_threads, num_non_epilogue_threads, num_epilogue_threads + }; + + if (get_env("DG_JIT_DEBUG") or get_env("DG_PRINT_CONFIGS")) { + const auto key = fmt::format( + "MegaMoESM90Config(num_ranks={}, num_experts={}, hidden={}, intermediate_hidden={}, num_max_tokens_per_rank={}, num_tokens={}, num_topk={})", + num_ranks, num_experts, hidden, intermediate_hidden, num_max_tokens_per_rank, num_tokens, num_topk); + static std::unordered_set printed; + if (printed.count(key) == 0) { + std::cout << key << ": " << config << std::endl; + printed.insert(key); + } + } + return config; +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp new file mode 100644 index 0000000000..a86ebe6f92 --- /dev/null +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -0,0 +1,305 @@ +#pragma once + +#include +#include "../../jit/compiler.hpp" +#include "../../jit/kernel_runtime.hpp" +#include "../../utils/exception.hpp" +#include "../../utils/format.hpp" +#include "runtime_utils.hpp" + +#include +#include + +#include "../heuristics/sm90_mega_moe.hpp" + +namespace deep_gemm { + +// ============================================================================ +// SM90 (Hopper) FP8 MegaMoE host runtime +// ---------------------------------------------------------------------------- +// This is the SM90 counterpart of `SM100FP8FP4MegaMoERuntime`. The kernel +// itself lives in `deep_gemm/impls/sm90_fp8_mega_moe.cuh`. +// +// Differences from SM100 path: +// * Activations and weights are both FP8 (e4m3); no FP4. +// * Activation/weight scale factors (SF) are float, not UE8M0 int + per-32 +// UTCCP layout. L1 activation SF and weight SF are per-128 K; the fused L1 +// epilogue writes L2 activation SF at per-64 K granularity. +// * No tensor memory: WGMMA accumulators are register-resident. +// * Cluster size is at most 2 (TMA multicast on A); no 2-CTA UMMA. +// ============================================================================ + +class SM90FP8MegaMoERuntime final : public LaunchRuntime { +public: + struct Args { + // Templated arguments + int num_max_tokens_per_rank; + int hidden, intermediate_hidden; + int num_experts, num_topk; + int num_ranks; + float activation_clamp; + bool fast_math; + int epilogue_registers; + bool reuse_accum_as_final; + bool l2_arrival_counter; + bool l2_epilogue_requires_full_sync; + bool split_phase_hot_path; + MegaMoESM90Config config; + + // Runtime arguments + void* y; + int* cumulative_local_expert_recv_stats; + int num_tokens; + layout::SymBuffer<> sym_buffer_ptrs; + + // Tensormaps for activations and weights. Weight scale factors use + // block (128, 128) quantization and are loaded by the math warpgroup + // directly from global memory (no TMA descriptor required). + CUtensorMap tensor_map_l1_acts; + CUtensorMap tensor_map_l1_acts_sf; + CUtensorMap tensor_map_l1_weights; + const float* l1_weights_sf; + CUtensorMap tensor_map_l1_output; + CUtensorMap tensor_map_l2_acts; + CUtensorMap tensor_map_l2_acts_sf; + CUtensorMap tensor_map_l2_weights; + const float* l2_weights_sf; + + // Launch configs + LaunchArgs launch_args; + }; + + static std::string generate_impl(const Args& args) { + return fmt::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_fp8_mega_moe_impl< + {}, + {}, {}, + {}, {}, + {}, + {}, {}, {}, + {}, + {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, + {} + >); +}}; +)", + args.num_max_tokens_per_rank, + args.hidden, args.intermediate_hidden, + args.num_experts, args.num_topk, + args.config.num_experts_per_wave, + args.config.block_m, args.config.block_n, args.config.block_k, + args.config.num_max_pool_tokens, + args.config.num_padded_sf_pool_tokens, + args.config.num_stages, + args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads, + args.launch_args.grid_dim.first, args.num_ranks, + to_string(args.activation_clamp), + args.fast_math ? "true" : "false", + args.epilogue_registers, + args.reuse_accum_as_final ? "true" : "false", + args.l2_arrival_counter ? "true" : "false", + args.l2_epilogue_requires_full_sync ? "true" : "false", + args.split_phase_hot_path ? "true" : "false"); + } + + static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { + DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + args.y, + args.cumulative_local_expert_recv_stats, + args.num_tokens, + args.sym_buffer_ptrs, + args.tensor_map_l1_acts, + args.tensor_map_l1_acts_sf, + args.tensor_map_l1_weights, + args.l1_weights_sf, + args.tensor_map_l1_output, + args.tensor_map_l2_acts, + args.tensor_map_l2_acts_sf, + args.tensor_map_l2_weights, + args.l2_weights_sf + )); + } +}; + +static void sm90_fp8_mega_moe( + const torch::Tensor& y, + const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf, + const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf, + const torch::Tensor& l1_weights, const torch::Tensor& l2_weights, + const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf, + const std::optional cumulative_local_expert_recv_stats, + const std::vector& sym_buffer_ptrs, + const int& rank_idx, const int& num_max_tokens_per_rank, + const int& num_experts_per_rank, + const int& num_tokens, const int& num_topk, + const int& hidden, const int& intermediate_hidden, + const float& activation_clamp, + const bool& fast_math +) { + const auto num_ranks = static_cast(sym_buffer_ptrs.size()); + const auto num_experts = num_experts_per_rank * num_ranks; + const auto num_padded_sf_pool_tokens = static_cast(l1_acts_sf.size(0)); + + // Heuristics + const auto config = get_mega_moe_config_sm90( + num_ranks, num_experts, num_experts_per_rank, + num_max_tokens_per_rank, num_tokens, num_topk, + hidden, intermediate_hidden, num_padded_sf_pool_tokens); + const int default_epilogue_registers = + config.num_epilogue_threads == 512 ? 112 : 0; + const int epilogue_registers = default_epilogue_registers; + if (epilogue_registers > 0) { + const int dispatch_registers = + config.num_epilogue_threads == 512 ? 32 : 48; + const int non_epilogue_registers = + config.num_epilogue_threads == 512 ? 24 : 40; + DG_HOST_ASSERT(dispatch_registers * config.num_dispatch_threads + + non_epilogue_registers * config.num_non_epilogue_threads + + epilogue_registers * config.num_epilogue_threads <= 64512); + } + const bool reuse_accum_as_final = config.block_m == 128; + const bool default_split_mn_barrier_opt = + config.block_m == 128 and config.block_n == 256 and + config.num_epilogue_threads == 512; + const bool split_phase_hot_path = + config.block_m == 128 and config.block_n == 256 and hidden >= 7168; + const bool decode_split_n_path = + config.block_m == 64 and config.num_epilogue_threads == 256; + const bool decode_split_n_bn256 = + decode_split_n_path and config.block_n == 256; + const bool decode_l2_counter = + decode_split_n_bn256 and num_tokens >= 4 and num_tokens <= 128; + const bool l2_arrival_counter = + default_split_mn_barrier_opt or decode_l2_counter; + const bool l2_epilogue_requires_full_sync = + not l2_arrival_counter; + + // Tensormap construction + // Acts/weights: standard 2D TMA descriptors (FP8 K-major). + // Activation SF: per-128 channel float for L1, per-64 for L2 (MN-major, no swizzle). + // Weight SF: block (128, 128) raw float pointer (no TMA descriptor). + constexpr int kGranK = 128; + constexpr int kL2ActsSFGranK = 64; + const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts, + hidden, config.num_max_pool_tokens, + config.block_k, config.block_m, + static_cast(l1_acts.stride(-2)), + config.swizzle_acts_mode); + const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf, + config.num_padded_sf_pool_tokens, hidden, + config.block_m, kGranK, + 1, 0); + const int weight_tma_block_n = config.block_n > 256 ? 256 : config.block_n; + const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights, + hidden, num_experts_per_rank * intermediate_hidden * 2, + config.block_k, weight_tma_block_n, + static_cast(l1_weights.stride(-2)), + config.swizzle_weights_mode); + // L1 output (post-SwiGLU FP8): N is halved. The correctness path stages + // this tile in plain row-major SMEM before the TMA store. Later L2 TMA + // loads may still swizzle from this row-major global buffer into their own + // SMEM tile. + // The usual TMA store is issued per warpgroup, each writing a `WG_BLOCK_M` + // row tile from its own SMEM offset. The m64n128 2-WG split-N decode path is + // different: both warpgroups stage one joint 64-column L1-output tile and a + // single warpgroup issues the combined store, so the descriptor must cover + // the full block_m x (block_n / 2) tile. + const int num_epilogue_warpgroups_h = config.num_epilogue_threads / 128; + const bool split_n_warpgroups = + config.block_m == 64 and num_epilogue_warpgroups_h > 1 and + config.block_n % num_epilogue_warpgroups_h == 0 and + (config.block_n / num_epilogue_warpgroups_h == 64 or + config.block_n / num_epilogue_warpgroups_h == 128); + const bool split_mn_warpgroups = + config.block_m == 128 and config.block_n == 256 and num_epilogue_warpgroups_h == 4; + const int wg_split_m = split_n_warpgroups ? 1 : + (split_mn_warpgroups ? 2 : num_epilogue_warpgroups_h); + const int wg_split_n = split_n_warpgroups ? num_epilogue_warpgroups_h : + (split_mn_warpgroups ? 2 : 1); + DG_HOST_ASSERT(wg_split_m * wg_split_n == num_epilogue_warpgroups_h); + const int wg_block_m = config.block_m / wg_split_m; + const int wg_block_n = config.block_n / wg_split_n; + const int wg_l1_out_block_n = wg_block_n / 2; + const bool split_n_shares_sf = + split_n_warpgroups and wg_l1_out_block_n < kL2ActsSFGranK; + const int l1_output_swizzle_mode = 0; + const int l1_output_box_n = + split_n_shares_sf ? config.block_n / 2 : wg_l1_out_block_n; + const int l1_output_box_m = + split_n_shares_sf ? config.block_m : wg_block_m; + const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts, + intermediate_hidden, config.num_max_pool_tokens, + l1_output_box_n, l1_output_box_m, + static_cast(l2_acts.stride(-2)), + l1_output_swizzle_mode); + const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts, + intermediate_hidden, config.num_max_pool_tokens, + config.block_k, config.block_m, + static_cast(l2_acts.stride(-2)), + config.swizzle_acts_mode); + const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf, + config.num_padded_sf_pool_tokens, intermediate_hidden, + config.block_m, kL2ActsSFGranK, + 1, 0); + const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights, + intermediate_hidden, num_experts_per_rank * hidden, + config.block_k, weight_tma_block_n, + static_cast(l2_weights.stride(-2)), + config.swizzle_weights_mode); + + // Stats can be optional + int* cumulative_local_expert_recv_stats_ptr = nullptr; + if (cumulative_local_expert_recv_stats.has_value()) + cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr(); + + // Launch + const auto num_sms = device_runtime->get_num_sms(); + const SM90FP8MegaMoERuntime::Args args = { + .num_max_tokens_per_rank = num_max_tokens_per_rank, + .hidden = hidden, .intermediate_hidden = intermediate_hidden, + .num_experts = num_experts, .num_topk = num_topk, + .num_ranks = num_ranks, + .activation_clamp = activation_clamp, + .fast_math = fast_math, + .epilogue_registers = epilogue_registers, + .reuse_accum_as_final = reuse_accum_as_final, + .l2_arrival_counter = l2_arrival_counter, + .l2_epilogue_requires_full_sync = l2_epilogue_requires_full_sync, + .split_phase_hot_path = split_phase_hot_path, + .config = config, + .y = y.data_ptr(), + .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, + .num_tokens = num_tokens, + .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + .tensor_map_l1_acts = tensor_map_l1_acts, + .tensor_map_l1_acts_sf = tensor_map_l1_acts_sf, + .tensor_map_l1_weights = tensor_map_l1_weights, + .l1_weights_sf = l1_weights_sf.data_ptr(), + .tensor_map_l1_output = tensor_map_l1_output, + .tensor_map_l2_acts = tensor_map_l2_acts, + .tensor_map_l2_acts_sf = tensor_map_l2_acts_sf, + .tensor_map_l2_weights = tensor_map_l2_weights, + .l2_weights_sf = l2_weights_sf.data_ptr(), + .launch_args = LaunchArgs(num_sms, config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, + config.smem_size, config.cluster_size) + }; + const auto code = SM90FP8MegaMoERuntime::generate(args); + const auto runtime = compiler->build("sm90_fp8_mega_moe", code); + SM90FP8MegaMoERuntime::launch(runtime, args); +} + +} // namespace deep_gemm diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index abd46179cd..4031aec1bd 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -18,6 +18,7 @@ #include "apis/gemm.hpp" #include "apis/layout.hpp" #include "apis/mega.hpp" +#include "apis/sm90_mega.hpp" #include "utils/torch_compat.hpp" @@ -613,6 +614,37 @@ dg_get_symm_buffer_size_for_mega_moe(int64_t num_ranks, int64_t num_experts, int num_bytes, slice_input_buffers); } +Tuple(TensorView)>> +dg_get_symm_buffer_size_for_sm90_mega_moe(int64_t num_ranks, int64_t num_experts, int64_t num_max_tokens_per_rank, int64_t num_topk, int64_t hidden, + int64_t intermediate_hidden, bool use_fp8_dispatch, std::string activation) { + auto [num_bytes, fn] = mega::get_symm_buffer_size_for_sm90_mega_moe( + static_cast(num_ranks), + static_cast(num_experts), + static_cast(num_max_tokens_per_rank), + static_cast(num_topk), + static_cast(hidden), + static_cast(intermediate_hidden), + use_fp8_dispatch, + activation + ); + + auto slice_input_buffers = [=](TensorView buffer) { + auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = fn(convert_to_torch_tensor(buffer)); + return Tuple( + Tensor::FromDLPack(at::toDLPack(x.view(at::kChar))), + Tensor::FromDLPack(at::toDLPack(x_sf)), + Tensor::FromDLPack(at::toDLPack(topk_idx)), + Tensor::FromDLPack(at::toDLPack(topk_weights)), + Tensor::FromDLPack(at::toDLPack(l1_acts.view(at::kChar))), + Tensor::FromDLPack(at::toDLPack(l1_acts_sf)), + Tensor::FromDLPack(at::toDLPack(l2_acts.view(at::kChar))), + Tensor::FromDLPack(at::toDLPack(l2_acts_sf)) + ); + }; + return Tuple(TensorView)>>( + num_bytes, slice_input_buffers); +} + void dg_fp8_fp4_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weights_sf, TensorView l2_weights, TensorView l2_weights_sf, Optional cumulative_local_expert_recv_stats, TensorView sym_buffer, Array sym_buffer_ptrs, int64_t rank_idx, int64_t num_max_tokens_per_rank, int64_t num_experts, int64_t num_topk, @@ -639,6 +671,32 @@ void dg_fp8_fp4_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weig } +void dg_fp8_mega_moe(TensorView y, TensorView l1_weights, TensorView l1_weights_sf, TensorView l2_weights, TensorView l2_weights_sf, + Optional cumulative_local_expert_recv_stats, TensorView sym_buffer, Array sym_buffer_ptrs, + int64_t rank_idx, int64_t num_max_tokens_per_rank, int64_t num_experts, int64_t num_topk, + Tuple recipe, std::string activation, Optional activation_clamp_opt, bool fast_math) { + auto c_val = cumulative_local_expert_recv_stats.has_value()? std::optional(convert_to_torch_tensor(cumulative_local_expert_recv_stats.value())) : std::nullopt; + auto act_clamp_opt_val = activation_clamp_opt.has_value()? std::optional(static_cast(activation_clamp_opt.value())) : std::nullopt; + std::vector sym_buffer_ptrs_val; + sym_buffer_ptrs_val.reserve(sym_buffer_ptrs.size()); + + for (Array::iterator it = sym_buffer_ptrs.begin(); it != sym_buffer_ptrs.end(); ++it) { + sym_buffer_ptrs_val.push_back(*it); + } + auto [recipe_a, recipe_b, recipe_c] = recipe; + auto recipe_val = std::make_tuple(static_cast(recipe_a), static_cast(recipe_b), static_cast(recipe_c)); + + mega::fp8_mega_moe( + convert_to_torch_tensor(y), + std::make_pair(convert_to_torch_tensor(l1_weights), convert_to_torch_tensor(l1_weights_sf)), + std::make_pair(convert_to_torch_tensor(l2_weights), convert_to_torch_tensor(l2_weights_sf)), + c_val, convert_to_torch_tensor(sym_buffer), sym_buffer_ptrs_val, static_cast(rank_idx), + static_cast(num_max_tokens_per_rank), static_cast(num_experts), + static_cast(num_topk), recipe_val, activation, act_clamp_opt_val, fast_math + ); +} + + void dg_mega_moe_pre_dispatch( TensorView x, TensorView topk_idx, TensorView topk_weights, TensorView buf_x, TensorView buf_x_sf, @@ -660,7 +718,9 @@ void dg_mega_moe_pre_dispatch( TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_token_alignment_for_mega_moe, dg_get_token_alignment_for_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_symm_buffer_size_for_mega_moe, dg_get_symm_buffer_size_for_mega_moe); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(get_symm_buffer_size_for_sm90_mega_moe, dg_get_symm_buffer_size_for_sm90_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_fp4_mega_moe, dg_fp8_fp4_mega_moe); +TVM_FFI_DLL_EXPORT_TYPED_FUNC(fp8_mega_moe, dg_fp8_mega_moe); TVM_FFI_DLL_EXPORT_TYPED_FUNC(mega_moe_pre_dispatch, dg_mega_moe_pre_dispatch); diff --git a/deep_gemm/include/deep_gemm/comm/barrier.cuh b/deep_gemm/include/deep_gemm/comm/barrier.cuh index 9c0fca6a42..3a40bb47ad 100644 --- a/deep_gemm/include/deep_gemm/comm/barrier.cuh +++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh @@ -67,9 +67,14 @@ CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, const auto start_clock = clock64(); while (ptx::ld_acq_sys(signal_ptr) != target) { if (clock64() - start_clock >= kNumTimeoutCycles) { +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) && (__CUDA_ARCH__ < 1000) && \ + !(defined(DG_NVLINK_BARRIER_VERBOSE_TIMEOUT) && DG_NVLINK_BARRIER_VERBOSE_TIMEOUT) + DG_TRAP_ONLY_DEVICE_ASSERT(false); +#else printf("DeepGEMM NVLink barrier timeout (180s): rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); DG_DEVICE_ASSERT(false and "NVLink barrier timeout"); +#endif } } } diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh new file mode 100644 index 0000000000..515f930257 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -0,0 +1,1935 @@ +#pragma once + +#pragma clang diagnostic push +#pragma clang diagnostic ignored "-Wunknown-attributes" + +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#define __CLION_IDE__ + +namespace deep_gemm { + +template +__forceinline__ __device__ float sm90_fp8_mega_moe_clamp_gate(float x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + return x; +} + +template +__forceinline__ __device__ float sm90_fp8_mega_moe_clamp_up(float x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + return x; +} + +template +__forceinline__ __device__ float sm90_fp8_mega_moe_silu(float x) { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; +} + +template +__forceinline__ __device__ float sm90_fp8_mega_moe_swiglu(float g, float u) { + g = sm90_fp8_mega_moe_clamp_gate(g); + u = sm90_fp8_mega_moe_clamp_up(u); + return sm90_fp8_mega_moe_silu(g) * u; +} + +__forceinline__ __device__ void sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( + const float2& amax, float2& sf, float2& sf_inv) { + constexpr float kScale = 1.0f / 448.0f; + const auto scaled = make_float2(__fmul_rn(amax.x, kScale), __fmul_rn(amax.y, kScale)); + const auto exp_x = math::fast_log2_ceil(scaled.x); + const auto exp_y = math::fast_log2_ceil(scaled.y); + sf.x = math::fast_pow2(exp_x), sf_inv.x = math::fast_pow2(-exp_x); + sf.y = math::fast_pow2(exp_y), sf_inv.y = math::fast_pow2(-exp_y); +} + +template +CUTLASS_DEVICE void sm90_fp8_mega_moe_for_each_block_split( + sched::MegaMoEScheduler& scheduler, + L1Func&& l1_func, L2Func&& l2_func) { + scheduler.fetch_expert_recv_count(); + scheduler.set_expert_idx(0); + + while (true) { + CUTE_TIE_DECL(scheduler.get_next_block(), block_phase, current_local_expert_idx, m_block_idx, n_block_idx); + if (block_phase == sched::BlockPhase::None) + break; + + if (block_phase == sched::BlockPhase::Linear1) { + l1_func(current_local_expert_idx, kNumL1BlockKs, m_block_idx, n_block_idx); + } else { + l2_func(current_local_expert_idx, kNumL2BlockKs, m_block_idx, n_block_idx); + } + } +} + +// ============================================================================ +// SM90 (Hopper) FP8 MegaMoE — full implementation +// ---------------------------------------------------------------------------- +// Pipeline (cluster=1, no TMA multicast): +// * Dispatch warps: pull tokens (FP8) and SF (per-128 channel float) from +// remote ranks via NVLink into the local L1 pool. +// * GEMM TMA-load warps (1 for A+SFA, 1 for B+SFB) feed the pipeline stages. +// * Math warpgroups (totalling kNumEpilogueThreads) consume each +// stage with WGMMA, accumulate into registers, then run the epilogue: +// - L1 (Linear1): SwiGLU with gate/up granularity-8 interleaved layout, +// per-row amax over each output-SF group, FP8 e4m3 quantize, STSM into +// SMEM, TMA store to local L1 output buffer. +// The per-row SF is written as a *float* into the L2-acts SF buffer at +// per-64 K granularity (one SF per L1 N block), so each block is fully +// self-contained and no cross-CTA amax synchronisation is needed. +// - L2 (Linear2): BF16 cast of the GEMM output, STSM into SMEM, then +// NVLink scatter to remote combine buffers. +// * After all GEMM blocks, the math warps run the COMBINE step (top-k +// reduction in BF16) — ported verbatim from the SM100 kernel. +// ============================================================================ + +template < + uint32_t kNumMaxTokensPerRank, + uint32_t kHidden, uint32_t kIntermediateHidden, + uint32_t kNumExperts, uint32_t kNumTopk, + uint32_t kNumExpertsPerWave, + uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K, + uint32_t kNumMaxPoolTokens, + uint32_t kNumPaddedSFPoolTokens, + uint32_t kNumStages, + uint32_t kNumDispatchThreads, uint32_t kNumNonEpilogueThreads, + uint32_t kNumEpilogueThreads, + uint32_t kNumSMs, uint32_t kNumRanks, + float kActivationClamp, + bool kFastMath, + uint32_t kEpilogueRegisterBudget, + bool kReuseAccumAsFinal, + bool kL2ArrivalCounter, + bool kL2EpilogueRequiresFullSync, + bool kSplitPhaseHotPath, + uint32_t L1_SHAPE_N = kIntermediateHidden * 2, + uint32_t L1_SHAPE_K = kHidden, + uint32_t L2_SHAPE_N = kHidden, + uint32_t L2_SHAPE_K = kIntermediateHidden, + uint32_t kNumDispatchWarps = kNumDispatchThreads / 32, + uint32_t kNumMMANonEpilogueWarps = kNumNonEpilogueThreads / 32, + uint32_t kNumEpilogueWarps = kNumEpilogueThreads / 32, + uint32_t kNumEpilogueWarpgroups = kNumEpilogueWarps / 4, + uint32_t kNumThreads = kNumDispatchThreads + kNumNonEpilogueThreads + kNumEpilogueThreads, + uint32_t kNumTokensPerWarp = 32 / kNumTopk, + uint32_t kNumExpertsPerRank = kNumExperts / kNumRanks +> +CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 1) void +sm90_fp8_mega_moe_impl(void* y, + int* cumulative_local_expert_recv_stats, + const uint32_t num_tokens, + const __grid_constant__ layout::SymBuffer sym_buffer, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_weights, + const float* __restrict__ l1_weights_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l1_output, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_acts_sf, + const __grid_constant__ cute::TmaDescriptor tensor_map_l2_weights, + const float* __restrict__ l2_weights_sf) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or defined(__CLION_IDE__) + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + // ===================================================================== + // Template checks + // ===================================================================== + DG_STATIC_ASSERT(kNumDispatchThreads >= 64 and kNumDispatchThreads % 64 == 0, + "Invalid number of dispatch threads"); + DG_STATIC_ASSERT(kNumNonEpilogueThreads == 64 or kNumNonEpilogueThreads == 128, + "Invalid number of GEMM TMA warps"); + DG_STATIC_ASSERT((kNumDispatchThreads + kNumNonEpilogueThreads) % 128 == 0, + "Math warpgroup start must be 128-thread aligned"); + DG_STATIC_ASSERT(kNumEpilogueThreads % 128 == 0, "Invalid number of math/epilogue threads"); + DG_STATIC_ASSERT(kNumExperts % kNumRanks == 0, "Invalid number of experts or ranks"); + DG_STATIC_ASSERT(BLOCK_M % 64 == 0, "BLOCK_M must be a multiple of WGMMA::M (64)"); + DG_STATIC_ASSERT(BLOCK_N == 128 or BLOCK_N == 256 or BLOCK_N == 512, + "SM90 MegaMoE supports CTA BLOCK_N=128/256/512"); + DG_STATIC_ASSERT(BLOCK_K == 128, "BLOCK_K is fixed to 128 (per-128 SF)"); + + // ===================================================================== + // Thread / warp identification + // ===================================================================== + const uint32_t sm_idx = blockIdx.x; + const uint32_t thread_idx = threadIdx.x; + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane_idx = ptx::get_lane_idx(); + + // Prefetch all TMA descriptors at the very beginning + if (warp_idx == 0 and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_l1_acts); + cute::prefetch_tma_descriptor(&tensor_map_l1_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l1_weights); + cute::prefetch_tma_descriptor(&tensor_map_l1_output); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts); + cute::prefetch_tma_descriptor(&tensor_map_l2_acts_sf); + cute::prefetch_tma_descriptor(&tensor_map_l2_weights); + } + + // ===================================================================== + // Workspaces and symmetric buffer slicing (mirror SM100 layout, except SF + // for L2 activations uses per-64 K granularity) + // ===================================================================== + const auto workspace = layout::Workspace( + sym_buffer.get_base_ptr(), kNumRanks, kNumExperts, kNumMaxTokensPerRank, kNumTopk); + + constexpr auto fp8_token_layout = layout::Data(kHidden); + constexpr auto bf16_token_layout = layout::Data(kHidden * sizeof(nv_bfloat16)); + constexpr auto fp8_intermediate_token_layout = layout::Data(kIntermediateHidden); + // Per-128 K float SF: 4 bytes per per-128 group => `kHidden / 32` bytes/token (same as SM100 packing) + constexpr auto fp8_sf_layout = layout::Data(kHidden / 32); + // Per-64 K float SF (SM90 only): 4 bytes per per-64 group => `kIntermediateHidden / 16` bytes/token + constexpr auto fp8_intermediate_sf_layout = layout::Data(kIntermediateHidden / 16); + constexpr auto input_topk_idx_layout = layout::Data(kNumTopk * sizeof(int64_t), false); + constexpr auto input_topk_weights_layout = layout::Data(kNumTopk * sizeof(float), false); + constexpr auto l1_topk_weights_layout = layout::Data(sizeof(float), false); + + // Registered input area + const auto input_token_buffer = layout::Buffer(fp8_token_layout, 1, kNumMaxTokensPerRank, workspace.get_end_ptr()); + const auto input_sf_buffer = layout::Buffer(fp8_sf_layout, 1, kNumMaxTokensPerRank, input_token_buffer.get_end_ptr()); + const auto input_topk_idx_buffer = layout::Buffer(input_topk_idx_layout, 1, kNumMaxTokensPerRank, input_sf_buffer.get_end_ptr()); + const auto input_topk_weights_buffer = layout::Buffer(input_topk_weights_layout, 1, kNumMaxTokensPerRank, input_topk_idx_buffer.get_end_ptr()); + + // L1 input area + const auto l1_token_buffer = layout::Buffer(fp8_token_layout, 1, kNumMaxPoolTokens, input_topk_weights_buffer.get_end_ptr()); + const auto l1_sf_buffer = layout::Buffer(fp8_sf_layout, 1, kNumPaddedSFPoolTokens, l1_token_buffer.get_end_ptr()); + const auto l1_topk_weights_buffer = layout::Buffer(l1_topk_weights_layout, 1, kNumMaxPoolTokens, l1_sf_buffer.get_end_ptr()); + + // L2 input area + const auto l2_token_buffer = layout::Buffer(fp8_intermediate_token_layout, 1, kNumMaxPoolTokens, l1_topk_weights_buffer.get_end_ptr()); + const auto l2_sf_buffer = layout::Buffer(fp8_intermediate_sf_layout, 1, kNumPaddedSFPoolTokens, l2_token_buffer.get_end_ptr()); + + // Combine input area + const auto combine_token_buffer = layout::Buffer(bf16_token_layout, kNumTopk, kNumMaxTokensPerRank, l2_sf_buffer.get_end_ptr()); + + // ===================================================================== + // GEMM data types and shape constants + // ===================================================================== + using a_dtype_t = cutlass::float_e4m3_t; + using b_dtype_t = cutlass::float_e4m3_t; + constexpr bool kSplitNWarpgroups = + BLOCK_M == 64 and kNumEpilogueWarpgroups > 1 and + BLOCK_N % kNumEpilogueWarpgroups == 0 and + ((BLOCK_N / kNumEpilogueWarpgroups == 64) or (BLOCK_N / kNumEpilogueWarpgroups == 128)); + constexpr bool kSplitMNWarpgroups = + BLOCK_M == 128 and BLOCK_N == 256 and kNumEpilogueWarpgroups == 4; + constexpr uint32_t kWarpgroupSplitM = kSplitNWarpgroups ? 1 : + (kSplitMNWarpgroups ? 2 : kNumEpilogueWarpgroups); + constexpr uint32_t kWarpgroupSplitN = kSplitNWarpgroups ? kNumEpilogueWarpgroups : + (kSplitMNWarpgroups ? 2 : 1); + constexpr uint32_t WG_BLOCK_M = BLOCK_M / kWarpgroupSplitM; + constexpr uint32_t WG_BLOCK_N = BLOCK_N / kWarpgroupSplitN; + constexpr uint32_t kNumCombineWarps = kNumEpilogueWarps; + using L1WGMMA = typename mma::sm90::FP8MMASelector::type; // M=64, N=WG_BLOCK_N, K=32 + using L2WGMMA = typename mma::sm90::FP8MMASelector::type; + constexpr uint32_t kL1OutputArrivalParts = 1; + static_assert(L1WGMMA::M == 64 and L1WGMMA::N == WG_BLOCK_N and L1WGMMA::K == 32, + "Unexpected WGMMA shape"); + DG_STATIC_ASSERT(kWarpgroupSplitM * kWarpgroupSplitN == kNumEpilogueWarpgroups, + "Invalid warpgroup split"); + DG_STATIC_ASSERT(WG_BLOCK_M == L1WGMMA::M, + "Each warpgroup must run exactly one WGMMA-M tile"); + DG_STATIC_ASSERT(kNumCombineWarps <= kNumEpilogueWarps, + "Combine warp count must fit in epilogue warps"); + + // Cluster=1 -> no multicast, A/B are loaded full-sized + constexpr uint32_t LOAD_BLOCK_M = BLOCK_M; + constexpr uint32_t LOAD_BLOCK_N = BLOCK_N; + constexpr uint32_t L1_OUT_BLOCK_N = BLOCK_N / 2; // post-SwiGLU + constexpr uint32_t WG_L1_OUT_BLOCK_N = WG_BLOCK_N / 2; + // When WG_L1_OUT_BLOCK_N < 64 the two N-split warpgroups jointly own a + // single per-64 L2-acts SF group, so they must publish ONE shared SF slot + // (k_sf_idx == n_block_idx) instead of one per warpgroup. The amax that + // feeds that shared SF must be reduced across both warpgroups. + constexpr bool kSplitNSharesSF = kSplitNWarpgroups and (WG_L1_OUT_BLOCK_N < 64); + constexpr uint32_t kSwizzleAMode = BLOCK_K * sizeof(a_dtype_t); // 128 + constexpr uint32_t kSwizzleBMode = BLOCK_K * sizeof(b_dtype_t); // 128 + constexpr uint32_t kSwizzleCDMode = 128; + constexpr uint32_t kGranK = 128; // L1 acts SF, weights SF + constexpr uint32_t kL2ActsSFGranK = 64; // L2 acts SF (per-64 K, SM90 only) + + // ===================================================================== + // Shared memory layout + // ===================================================================== + constexpr uint32_t kSharedMemoryAlignment = 1024; + extern __shared__ __align__(kSharedMemoryAlignment) uint8_t smem_buffer[]; + + constexpr uint32_t SMEM_EXPERT_COUNT_SIZE = + math::constexpr_align(kNumExperts * sizeof(uint32_t), kSharedMemoryAlignment); + constexpr uint32_t SMEM_SEND_BUFFER_SIZE = + math::constexpr_align(fp8_token_layout.get_num_bytes() * kNumDispatchWarps, kSharedMemoryAlignment); + constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); + constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); + // SFA per-stage must be sized for the larger of L1 (BLOCK_M floats) and L2 (2*BLOCK_M floats per-64). + constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = + math::constexpr_align(2 * BLOCK_M * sizeof(float), 128u); + // Block (128, 128) weight SF is loaded directly from global by the math + // warpgroup, so no SMEM is needed. + constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = 0; + + // CD output: max of L1 FP8 (BLOCK_M * (BLOCK_N/2) * 1 byte) and + // L2 BF16 (BLOCK_M * BLOCK_N * 2 bytes). Split-M warpgroups own disjoint + // row slices; shared-SF split-N warpgroups stage disjoint column slices + // into one CTA tile. + constexpr uint32_t SMEM_CD_L1_SIZE = BLOCK_M * L1_OUT_BLOCK_N * sizeof(cutlass::float_e4m3_t); + constexpr uint32_t SMEM_CD_L2_SIZE = BLOCK_M * BLOCK_N * sizeof(nv_bfloat16); + constexpr uint32_t SMEM_CD_SIZE = math::constexpr_align( + SMEM_CD_L1_SIZE > SMEM_CD_L2_SIZE ? SMEM_CD_L1_SIZE : SMEM_CD_L2_SIZE, kSharedMemoryAlignment); + + constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE = + SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_CD_SIZE + + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + + // SMEM pointers + auto smem_expert_count = reinterpret_cast(smem_buffer); + const auto smem_send_buffers = layout::Buffer( + fp8_token_layout, kNumDispatchWarps, 1, + math::advance_ptr(smem_buffer, SMEM_EXPERT_COUNT_SIZE)); + + auto smem_gemm_base = math::advance_ptr( + smem_buffer, SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE); + + // CD output is shared by L1 (FP8) and L2 (BF16); reinterpret-cast as needed. + auto smem_cd_l1 = reinterpret_cast(smem_gemm_base); + auto smem_cd_l2 = reinterpret_cast(smem_gemm_base); + + auto smem_a = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); + }); + auto smem_b = utils::PatternVisitor([=](const uint32_t& i) { + return math::advance_ptr(smem_gemm_base, SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); + }); + auto sf_start_ptr = math::advance_ptr(smem_gemm_base, + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); + auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { + return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); + }); + + // Barriers live after SF (SFB is loaded directly from global, no SMEM) + auto barrier_start_ptr = reinterpret_cast( + sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE); + auto dispatch_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + i; }); + auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + i; }); + auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages + i; }); + auto combine_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + kNumDispatchWarps + kNumStages * 2 + i; }); + + // ===================================================================== + // Initialization + // ===================================================================== + if (warp_idx == 0) { + // Clean expert-count shared memory + #pragma unroll + for (uint32_t i = lane_idx; i < kNumExperts; i += 32) + ptx::st_shared(smem_expert_count + i, 0u); + } else if (warp_idx == 1) { + // Init dispatch m-barriers + #pragma unroll + for (uint32_t i = lane_idx; i < kNumDispatchWarps; i += 32) + dispatch_barriers[i]->init(1); + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 2) { + // Init GEMM full/empty barriers and combine barriers + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t i = 0; i < kNumStages; ++ i) { + // Two producer warps (A+SFA loader, B+SFB loader) each call + // `arrive_and_expect_tx` per stage, so init count must be 2. + full_barriers[i]->init(2); + // Each math warp arrives once per stage release. + empty_barriers[i]->init(kNumEpilogueWarps); + } + #pragma unroll + for (uint32_t i = 0; i < kNumCombineWarps * 2; ++ i) + combine_barriers[i]->init(1); + } + cutlass::arch::fence_barrier_init(); + } + __syncthreads(); + + // ===================================================================== + // Scheduler (cluster=1) + // ===================================================================== + auto scheduler = sched::MegaMoEScheduler< + BLOCK_M, BLOCK_N, BLOCK_K, + L1_SHAPE_N, L1_SHAPE_K, + L2_SHAPE_N, L2_SHAPE_K, + kNumExpertsPerRank, kNumExpertsPerWave, + kNumSMs, kNumRanks>(workspace); + + // Pipeline state shared by TMA loaders and math warpgroups + uint32_t stage_idx = 0, phase = 0; + auto advance_pipeline = [&](uint32_t& k_block_idx) { + ++ k_block_idx; + stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1; + phase ^= stage_idx == 0; + }; + + // Intra-SM barrier indices (mirroring SM100) + constexpr uint32_t kDispatchBarrierIdx = 0; + constexpr uint32_t kDispatchWithEpilogueBarrierIdx = 1; + constexpr uint32_t kEpilogueFullBarrierIdx = 2; + constexpr uint32_t kEpilogueWGBarrierStartIdx = 3; + + // Cross-rank NVLink barrier tags + constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; + constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; + constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + + // Register reconfiguration counts (chosen to fit in 64512 reg budget). + // For the 256-epilogue-thread split-N decode path: + // 64*48 + 64*40 + 256*168 = 48640 <= 64512. + // For the 512-epilogue-thread split-MN path, trim dispatch and loader roles + // so launch bounds still leave enough WGMMA registers. + // Reduced-thread decode (kNumThreads<=256) raises the launch-bounds + // register ceiling to 65536/256=256; grant the epilogue warpgroup the full + // 256 so the accumulator double-buffer fits without spilling. + // 64*48 + 64*40 + 128*256 = 38400 <= 64512. + constexpr uint32_t kNumEpilogueRegisters = + kEpilogueRegisterBudget == 0 ? + (kNumEpilogueThreads == 512 ? 112 : + (kNumEpilogueThreads == 256 ? 168 : + (kNumThreads <= 256u ? 256 : 208))) : + kEpilogueRegisterBudget; + // The 512-epilogue-thread path has only 3584 registers of headroom at + // epilogue=112. Raising epilogue to 120 is not viable without changing + // the role topology: dispatch=24 stalls the split-MN path, while + // non-epilogue=16 is below ptxas' setmaxnreg.dec legal minimum. + constexpr uint32_t kNumDispatchRegisters = + kNumEpilogueThreads == 512 ? 32 : 48; + constexpr uint32_t kNumNonEpilogueRegisters = + kNumEpilogueThreads == 512 ? 24 : 40; + DG_STATIC_ASSERT(kNumDispatchRegisters * kNumDispatchThreads + + kNumNonEpilogueRegisters * kNumNonEpilogueThreads + + kNumEpilogueRegisters * kNumEpilogueThreads <= 64512, + "Too many registers"); + + constexpr uint32_t kDispatchGridSyncIndex = 0; + constexpr uint32_t kEpilogueGridSyncIndex = 1; + + // ===================================================================== + // ROLE 1: DISPATCH WARPS + // Mirrors SM100 dispatch with two changes: + // * SF is per-128 channel float (no UTCCP transpose). We store the + // remote per-token SF directly into the local L1 SF buffer in + // MN-major layout: `local_sf[k_chunk * num_padded_sf_pool_tokens + token_idx]`. + // * The "token_idx_in_expert" → SF token index is now the simple + // per-block linear mapping (no 4×32 transpose). + // ===================================================================== + if (warp_idx < kNumDispatchWarps) { + cutlass::arch::warpgroup_reg_dealloc(); + + DG_STATIC_ASSERT(kNumTopk <= 32, "Invalid number of topk"); + constexpr uint32_t kNumActivateLanes = kNumTokensPerWarp * kNumTopk; + const auto read_topk_idx = [&](const auto& process) { + #pragma unroll + for (uint32_t i = (sm_idx * kNumDispatchWarps + warp_idx) * kNumTokensPerWarp; + i < num_tokens; + i += kNumSMs * kNumDispatchWarps * kNumTokensPerWarp) { + int expert_idx = -1; + if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { + expert_idx = static_cast( + __ldg(input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + if (expert_idx >= 0) + process(i * kNumTopk + lane_idx, expert_idx); + } + __syncwarp(); + } + }; + + // Count tokens per expert + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + atomicAdd_block(smem_expert_count + expert_idx, 1); + }); + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Stake out per-expert SM offsets via global atomic + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const uint64_t send_value = (1ull << 32) | static_cast(smem_expert_count[i]); + smem_expert_count[i] = static_cast( + ptx::atomic_add(workspace.get_expert_send_count_ptr(i), send_value)); + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + // Write source token-topk indices to remote ranks + read_topk_idx([&](const uint32_t& token_topk_idx, const int& expert_idx) { + const auto dst_rank_idx = expert_idx / kNumExpertsPerRank; + const auto dst_slot_idx = atomicAdd_block(smem_expert_count + expert_idx, 1); + const auto dst_ptr = workspace.get_src_token_topk_idx_ptr( + expert_idx % kNumExpertsPerRank, sym_buffer.rank_idx, dst_slot_idx); + *sym_buffer.map(dst_ptr, dst_rank_idx) = token_topk_idx; + }); + + comm::grid_sync( + workspace, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); } + ); + + if (sm_idx == 0) { + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { + const auto dst_rank_idx = i / kNumExpertsPerRank; + const auto dst_local_expert_idx = i % kNumExpertsPerRank; + const auto expert_status = *workspace.get_expert_send_count_ptr(i); + *sym_buffer.map( + workspace.get_expert_recv_count_ptr(sym_buffer.rank_idx, dst_local_expert_idx), + dst_rank_idx) = expert_status & 0xffffffff; + ptx::atomic_add_sys( + sym_buffer.map(workspace.get_expert_recv_count_sum_ptr(dst_local_expert_idx), dst_rank_idx), + expert_status); + } + } + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + false, true); + + // Sync with epilogue warps before pulling tokens. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Token / SF pull loop + uint32_t pull_mbarrier_phase = 0; + const auto pull_buffer = smem_send_buffers.get_rank_buffer(warp_idx).get_data_buffer(0); + const auto pull_mbarrier = dispatch_barriers[warp_idx]; + + scheduler.fetch_expert_recv_count(); + + constexpr uint32_t kNumRanksPerLane = math::constexpr_ceil_div(kNumRanks, 32u); + int current_expert_idx = -1; + uint32_t stored_rank_count[kNumRanksPerLane] = {}; + uint32_t expert_start_idx = 0, expert_end_idx = 0; + uint32_t expert_pool_block_offset = 0; + + constexpr uint32_t kNumGlobalWarps = kNumSMs * kNumDispatchWarps; + for (uint32_t token_idx = sm_idx * kNumDispatchWarps + warp_idx; ; token_idx += kNumGlobalWarps) { + int old_expert_idx = current_expert_idx; + while (token_idx >= expert_end_idx) { + if (++ current_expert_idx >= kNumExpertsPerRank) + break; + expert_pool_block_offset += math::ceil_div(expert_end_idx - expert_start_idx, BLOCK_M); + expert_start_idx = expert_end_idx; + expert_end_idx += scheduler.get_num_tokens(current_expert_idx); + } + if (current_expert_idx >= kNumExpertsPerRank) + break; + + if (old_expert_idx != current_expert_idx) { + old_expert_idx = current_expert_idx; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t j = i * 32 + lane_idx; + stored_rank_count[i] = j < kNumRanks ? + static_cast(*workspace.get_expert_recv_count_ptr(j, current_expert_idx)) : 0; + } + } + + // Round-robin rank selection (identical to SM100) + uint32_t current_rank_in_expert_idx; + uint32_t remaining[kNumRanksPerLane]; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] = stored_rank_count[i]; + uint32_t offset = 0; + uint32_t token_idx_in_expert = token_idx - expert_start_idx; + uint32_t slot_idx = token_idx_in_expert; + uint32_t token_idx_in_rank; + while (true) { + uint32_t num_actives_in_lane = 0; + uint32_t min_in_lane = 0xffffffff; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + num_actives_in_lane += remaining[i] > 0; + if (remaining[i] > 0) + min_in_lane = cute::min(min_in_lane, remaining[i]); + } + const uint32_t num_active_ranks = __reduce_add_sync(0xffffffff, num_actives_in_lane); + const uint32_t length = __reduce_min_sync(0xffffffff, min_in_lane); + + const uint32_t num_round_tokens = length * num_active_ranks; + if (slot_idx < num_round_tokens) { + const uint32_t slot_idx_in_round = slot_idx % num_active_ranks; + uint32_t num_seen_ranks = 0; + current_rank_in_expert_idx = 0; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) { + const uint32_t mask = __ballot_sync(0xffffffff, remaining[i] > 0); + const uint32_t num_active_lanes = __popc(mask); + if (slot_idx_in_round >= num_seen_ranks and slot_idx_in_round < num_seen_ranks + num_active_lanes) + current_rank_in_expert_idx = i * 32 + __fns(mask, 0, slot_idx_in_round - num_seen_ranks + 1); + num_seen_ranks += num_active_lanes; + } + token_idx_in_rank = offset + (slot_idx / num_active_ranks); + break; + } + slot_idx -= num_round_tokens; + offset += length; + #pragma unroll + for (uint32_t i = 0; i < kNumRanksPerLane; ++ i) + remaining[i] -= cute::min(remaining[i], length); + } + + const uint32_t src_token_topk_idx = *workspace.get_src_token_topk_idx_ptr( + current_expert_idx, current_rank_in_expert_idx, token_idx_in_rank); + const uint32_t src_token_idx = src_token_topk_idx / kNumTopk; + const uint32_t src_topk_idx = src_token_topk_idx % kNumTopk; + + const uint32_t pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; + + // Pull token data. Overlap a remote TMA load with SF copy and + // then use TMA store to materialize the local L1 input. + if (cute::elect_one_sync()) { + ptx::tma_load_1d( + pull_buffer.get_base_ptr(), + sym_buffer.map(input_token_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx), + pull_mbarrier, kHidden); + } + __syncwarp(); + + // Copy SF: per-128 K floats, written linearly (no UTCCP transpose). + constexpr uint32_t kNumSFFloats = kHidden / 128; + DG_STATIC_ASSERT(kNumSFFloats > 0 and kHidden % 128 == 0, "Invalid SF"); + const auto remote_sf_ptr = sym_buffer.map( + input_sf_buffer.get_data_buffer(src_token_idx).get_base_ptr(), + current_rank_in_expert_idx); + const auto local_sf_ptr = l1_sf_buffer.get_base_ptr(); + const uint32_t sf_pool_token_idx = expert_pool_block_offset * BLOCK_M + token_idx_in_expert; + #pragma unroll + for (uint32_t i = 0; i < math::constexpr_ceil_div(kNumSFFloats, 32u); ++ i) { + const uint32_t j = i * 32 + lane_idx; + if (j < kNumSFFloats) + local_sf_ptr[j * kNumPaddedSFPoolTokens + sf_pool_token_idx] = remote_sf_ptr[j]; + } + __syncwarp(); + + if (cute::elect_one_sync()) { + const auto weight = *sym_buffer.map( + input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + *l1_topk_weights_buffer.get_data_buffer(pool_token_idx).get_base_ptr() = weight; + } + __syncwarp(); + + if (cute::elect_one_sync()) { + ptx::mbarrier_arrive_and_set_tx(pull_mbarrier, kHidden); + ptx::mbarrier_wait_and_flip_phase(pull_mbarrier, pull_mbarrier_phase); + + ptx::tma_store_1d( + l1_token_buffer.get_data_buffer(pool_token_idx).get_base_ptr(), + pull_buffer.get_base_ptr(), pull_buffer.get_num_bytes()); + + *workspace.get_token_src_metadata_ptr(pool_token_idx) = + {current_rank_in_expert_idx, src_token_idx, src_topk_idx}; + + cute::tma_store_arrive(); + ptx::tma_store_wait<0>(); + ptx::red_add_rel( + workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + token_idx_in_expert / BLOCK_M), 1); + } + __syncwarp(); + } + + // Cleanup workspace, overlapping with combine. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + DG_STATIC_ASSERT(kNumSMs > 1, "Invalid SM count"); + if (sm_idx == 0) { + #pragma unroll + for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) + *workspace.get_expert_send_count_ptr(i) = 0; + } else { + for (uint32_t i = sm_idx - 1; i < kNumExpertsPerRank; i += kNumSMs - 1) { + const auto num_recv_tokens = static_cast( + *workspace.get_expert_recv_count_sum_ptr(i)); + const auto num_recv_m_blocks = math::ceil_div(num_recv_tokens, BLOCK_M); + + expert_pool_block_offset = scheduler.get_pool_block_offset(i); + + ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); + + DG_STATIC_ASSERT(kNumDispatchWarps >= 2, "Not enough dispatch warps"); + if (warp_idx == 0) { + *workspace.get_expert_recv_count_sum_ptr(i) = 0; + } else if (warp_idx == 1) { + if (cute::elect_one_sync() and cumulative_local_expert_recv_stats != nullptr) + ptx::red_add(cumulative_local_expert_recv_stats + i, static_cast(num_recv_tokens)); + __syncwarp(); + } + + for (uint32_t j = thread_idx; j < kNumRanks; j += kNumDispatchThreads) + *workspace.get_expert_recv_count_ptr(j, i) = 0; + __syncwarp(); + + for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { + *workspace.get_l1_arrival_count_ptr(expert_pool_block_offset + j) = 0; + *workspace.get_l2_arrival_mask_ptr(expert_pool_block_offset + j) = 0; + } + __syncwarp(); + } + } + + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, thread_idx, + [=]() { ptx::sync_aligned(kNumDispatchThreads, kDispatchBarrierIdx); }, + true, false); + + // ===================================================================== + // ROLE 2: GEMM TMA LOAD warps (load A+SFA, B+SFB) + // Warps inside `kNumNonEpilogueThreads`: warp 0 loads A + SFA, + // warp 1 loads B. + // ===================================================================== + } else if (warp_idx == kNumDispatchWarps) { + cutlass::arch::warpgroup_reg_dealloc(); + + auto process_a_sfa_block = [&](const auto& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + const auto tensor_map_a_ptr = block_phase == sched::BlockPhase::Linear2 + ? &tensor_map_l2_acts : &tensor_map_l1_acts; + const auto tensor_map_sfa_ptr = block_phase == sched::BlockPhase::Linear2 + ? &tensor_map_l2_acts_sf : &tensor_map_l1_acts_sf; + + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + + // Wait for the pool to be ready + if (block_phase == sched::BlockPhase::Linear1) { + const auto ptr = workspace.get_l1_arrival_count_ptr(pool_block_idx); + const auto expected = scheduler.template get_valid_m(); + while (ptx::ld_acq(ptr) != expected); + } else { + constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; + if constexpr (kL2ArrivalCounter) { + const auto ptr = reinterpret_cast( + workspace.get_l2_arrival_mask_ptr(pool_block_idx)); + const uint32_t active_m_wgs = math::ceil_div( + scheduler.template get_valid_m(), WG_BLOCK_M); + const uint32_t expected = + kNumL1BlockNs * active_m_wgs * kWarpgroupSplitN * kL1OutputArrivalParts; + while (ptx::ld_acq(ptr) != expected); + } else { + const auto ptr = workspace.get_l2_arrival_mask_ptr(pool_block_idx); + // Each L1 N block sets one bit; total bits = L1_SHAPE_N / BLOCK_N. + const uint64_t expected = (kNumL1BlockNs >= 64) + ? ~0ull : ((1ull << kNumL1BlockNs) - 1ull); + while (ptx::ld_acq_gpu(ptr) != expected); + } + } + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + if (cute::elect_one_sync()) { + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // TMA load A + tma::copy( + tensor_map_a_ptr, full_barriers[stage_idx], smem_a[stage_idx], + k_idx, m_idx, 1); + + // TMA load SFA + if (block_phase == sched::BlockPhase::Linear1) { + // L1 SFA per-128: load (BLOCK_M, 1) at K=k_block_idx + tma::copy( + tensor_map_sfa_ptr, full_barriers[stage_idx], smem_sfa[stage_idx], + m_idx, k_block_idx, 1); + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + BLOCK_M * sizeof(float)); + } else { + // L2 SFA per-64: descriptor box is (block_mn, 1) (see make_tma_sf_desc), + // so we must issue two single-group TMAs and place them at smem offsets + // 0 and BLOCK_M to match math's load offsets (`+ 0 * BLOCK_M` / `+ 1 * BLOCK_M`). + tma::copy( + tensor_map_sfa_ptr, full_barriers[stage_idx], smem_sfa[stage_idx], + m_idx, k_block_idx * 2, 1); + tma::copy( + tensor_map_sfa_ptr, full_barriers[stage_idx], + smem_sfa[stage_idx] + BLOCK_M, + m_idx, k_block_idx * 2 + 1, 1); + full_barriers[stage_idx]->arrive_and_expect_tx( + SMEM_A_SIZE_PER_STAGE + 2 * BLOCK_M * sizeof(float)); + } + } + __syncwarp(); + } + }; + + if constexpr (kSplitPhaseHotPath) { + sm90_fp8_mega_moe_for_each_block_split( + scheduler, + [&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_a_sfa_block( + std::integral_constant{}, + local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }, + [&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_a_sfa_block( + std::integral_constant{}, + local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }); + } else { + scheduler.for_each_block([&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_a_sfa_block(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }); + } + + } else if (warp_idx == kNumDispatchWarps + 1) { + cutlass::arch::warpgroup_reg_dealloc(); + + scheduler.for_each_block([&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + const auto tensor_map_b_ptr = + block_phase == sched::BlockPhase::Linear2 ? &tensor_map_l2_weights : &tensor_map_l1_weights; + + const uint32_t shape_n = block_phase == sched::BlockPhase::Linear2 ? L2_SHAPE_N : L1_SHAPE_N; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + empty_barriers[stage_idx]->wait(phase ^ 1); + + if (cute::elect_one_sync()) { + const uint32_t n_idx = local_expert_idx * shape_n + n_block_idx * BLOCK_N; + const uint32_t k_idx = k_block_idx * BLOCK_K; + + // TMA load B (weight SF is now loaded directly by math warps from global) + if constexpr (LOAD_BLOCK_N <= 256) { + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], smem_b[stage_idx], + k_idx, n_idx, 1); + } else { + DG_STATIC_ASSERT(LOAD_BLOCK_N % 256 == 0, + "Large B tiles are loaded as 256-column TMA slices"); + #pragma unroll + for (uint32_t b_slice_idx = 0; b_slice_idx < LOAD_BLOCK_N / 256; ++ b_slice_idx) { + tma::copy( + tensor_map_b_ptr, full_barriers[stage_idx], + smem_b[stage_idx] + b_slice_idx * 256 * BLOCK_K, + k_idx, n_idx + b_slice_idx * 256, 1); + } + } + + full_barriers[stage_idx]->arrive_and_expect_tx(SMEM_B_SIZE_PER_STAGE); + } + __syncwarp(); + } + }); + + } else if (warp_idx < kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // Idle non-epilogue warps (kNumDispatchWarps+2, +3). They must still + // participate in the warpgroup-collective `setmaxnreg.dec.sync.aligned` + // so that the math warpgroup's `warpgroup_reg_alloc` can succeed. + cutlass::arch::warpgroup_reg_dealloc(); + + } else if (warp_idx >= kNumDispatchWarps + kNumMMANonEpilogueWarps) { + // ===================================================================== + // ROLE 3: MATH WARPGROUPS (WGMMA + epilogue + combine) + // ===================================================================== + cutlass::arch::warpgroup_reg_alloc(); + + const uint32_t epilogue_warp_idx = warp_idx - (kNumDispatchWarps + kNumMMANonEpilogueWarps); + const uint32_t epilogue_wg_idx = epilogue_warp_idx / 4; + const uint32_t epilogue_thread_idx = epilogue_warp_idx * 32 + lane_idx; + const uint32_t warp_idx_in_wg = epilogue_warp_idx % 4; + + // WGMMA-output register layout helpers + const uint32_t row_idx = lane_idx / 4; + const uint32_t col_idx = lane_idx % 4; + const uint32_t r_0 = warp_idx_in_wg * 16 + row_idx; + const uint32_t r_1 = r_0 + 8; + + // When the two N-split warpgroups share a single per-64 SF group they + // also stage into ONE shared row-major L1-output tile (stride + // L1_OUT_BLOCK_N), each writing its own WG_L1_OUT_BLOCK_N-column half, + // so a single combined TMA store matches the host descriptor box. + constexpr uint32_t WG_SMEM_CD_L1_STRIDE_N = + kSplitNSharesSF ? L1_OUT_BLOCK_N : WG_L1_OUT_BLOCK_N; + constexpr uint32_t WG_SMEM_CD_L2_STRIDE_N = WG_BLOCK_N; + + // Sync with dispatch in the full communication path. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + auto process_math_block = [&](const auto& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + const uint32_t valid_m = scheduler.template get_valid_m(); + const uint32_t pool_block_idx = scheduler.get_current_pool_block_offset() + m_block_idx; + const uint32_t m_idx = pool_block_idx * BLOCK_M; + const uint32_t n_idx = n_block_idx * BLOCK_N; + const uint32_t epilogue_wg_m_idx = epilogue_wg_idx / kWarpgroupSplitN; + const uint32_t epilogue_wg_n_idx = epilogue_wg_idx - epilogue_wg_m_idx * kWarpgroupSplitN; + const uint32_t wg_n_offset = epilogue_wg_n_idx * WG_BLOCK_N; + const uint32_t wg_l1_out_n_offset = epilogue_wg_n_idx * WG_L1_OUT_BLOCK_N; + const uint32_t row_base = epilogue_wg_m_idx * WG_BLOCK_M; + const uint32_t row_offset_r0 = row_base + r_0; + const uint32_t row_offset_r1 = row_base + r_1; + const uint32_t sf_n_block_idx = kSplitNSharesSF ? n_block_idx + : (n_block_idx * kWarpgroupSplitN + epilogue_wg_n_idx); + const uint32_t smem_a_wg_offset = epilogue_wg_m_idx * WG_BLOCK_M * BLOCK_K; + const uint32_t smem_b_wg_offset = epilogue_wg_n_idx * WG_BLOCK_N * BLOCK_K; + // In the shared-tile case the WG stages into the joint L1-output tile + // at its own column offset (row stride L1_OUT_BLOCK_N); otherwise each + // WG owns a disjoint contiguous WG_BLOCK_M x WG_L1_OUT_BLOCK_N slice. + const uint32_t smem_cd_l1_wg_offset = kSplitNSharesSF ? wg_l1_out_n_offset + : (epilogue_wg_idx * WG_BLOCK_M * WG_L1_OUT_BLOCK_N); + const uint32_t smem_cd_l2_wg_offset = epilogue_wg_idx * WG_BLOCK_M * WG_BLOCK_N; + const bool valid_r0 = row_offset_r0 < valid_m; + const bool valid_r1 = row_offset_r1 < valid_m; + + // ---------------- GEMM ---------------- + using WGMMA = L1WGMMA; + constexpr uint32_t kAccumPerThread = WGMMA::kNumAccum; + float final_accum[kAccumPerThread] = {}; + + if constexpr (kReuseAccumAsFinal) { + auto prescale_l1_final = [&](const float& scale_a_0, const float& scale_a_1, + const float& gate_sf, const float& up_sf) { + const float inv_s0_gate = kFastMath ? math::fast_rcp(scale_a_0 * gate_sf) : 1.0f / (scale_a_0 * gate_sf); + const float inv_s1_gate = kFastMath ? math::fast_rcp(scale_a_1 * gate_sf) : 1.0f / (scale_a_1 * gate_sf); + const float inv_s0_up = kFastMath ? math::fast_rcp(scale_a_0 * up_sf) : 1.0f / (scale_a_0 * up_sf); + const float inv_s1_up = kFastMath ? math::fast_rcp(scale_a_1 * up_sf) : 1.0f / (scale_a_1 * up_sf); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float inv_s0 = (i & 1u) ? inv_s0_up : inv_s0_gate; + const float inv_s1 = (i & 1u) ? inv_s1_up : inv_s1_gate; + final_accum[i*4+0] *= inv_s0; + final_accum[i*4+1] *= inv_s0; + final_accum[i*4+2] *= inv_s1; + final_accum[i*4+3] *= inv_s1; + } + }; + auto postscale_l1_final = [&](const float& scale_a_0, const float& scale_a_1, + const float& gate_sf, const float& up_sf) { + const float s0_gate = scale_a_0 * gate_sf; + const float s1_gate = scale_a_1 * gate_sf; + const float s0_up = scale_a_0 * up_sf; + const float s1_up = scale_a_1 * up_sf; + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float s0 = (i & 1u) ? s0_up : s0_gate; + const float s1 = (i & 1u) ? s1_up : s1_gate; + final_accum[i*4+0] *= s0; + final_accum[i*4+1] *= s0; + final_accum[i*4+2] *= s1; + final_accum[i*4+3] *= s1; + } + }; + auto prescale_l2_final = [&](const float& scale_a_0, const float& scale_a_1, + const float& l2_sf) { + const float inv_s0 = kFastMath ? math::fast_rcp(scale_a_0 * l2_sf) : 1.0f / (scale_a_0 * l2_sf); + const float inv_s1 = kFastMath ? math::fast_rcp(scale_a_1 * l2_sf) : 1.0f / (scale_a_1 * l2_sf); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] *= inv_s0; + final_accum[i*4+1] *= inv_s0; + final_accum[i*4+2] *= inv_s1; + final_accum[i*4+3] *= inv_s1; + } + }; + auto postscale_l2_final = [&](const float& scale_a_0, const float& scale_a_1, + const float& l2_sf) { + const float s0 = scale_a_0 * l2_sf; + const float s1 = scale_a_1 * l2_sf; + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] *= s0; + final_accum[i*4+1] *= s0; + final_accum[i*4+2] *= s1; + final_accum[i*4+3] *= s1; + } + }; + auto rescale_l1_final = [&](const float& prev_scale_a_0, const float& prev_scale_a_1, + const float& prev_gate_sf, const float& prev_up_sf, + const float& scale_a_0, const float& scale_a_1, + const float& gate_sf, const float& up_sf) { + const float r0_gate = (prev_scale_a_0 * prev_gate_sf) * + (kFastMath ? math::fast_rcp(scale_a_0 * gate_sf) : 1.0f / (scale_a_0 * gate_sf)); + const float r1_gate = (prev_scale_a_1 * prev_gate_sf) * + (kFastMath ? math::fast_rcp(scale_a_1 * gate_sf) : 1.0f / (scale_a_1 * gate_sf)); + const float r0_up = (prev_scale_a_0 * prev_up_sf) * + (kFastMath ? math::fast_rcp(scale_a_0 * up_sf) : 1.0f / (scale_a_0 * up_sf)); + const float r1_up = (prev_scale_a_1 * prev_up_sf) * + (kFastMath ? math::fast_rcp(scale_a_1 * up_sf) : 1.0f / (scale_a_1 * up_sf)); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float r0 = (i & 1u) ? r0_up : r0_gate; + const float r1 = (i & 1u) ? r1_up : r1_gate; + final_accum[i*4+0] *= r0; + final_accum[i*4+1] *= r0; + final_accum[i*4+2] *= r1; + final_accum[i*4+3] *= r1; + } + }; + auto rescale_l2_final = [&](const float& prev_scale_a_0, const float& prev_scale_a_1, + const float& prev_l2_sf, + const float& scale_a_0, const float& scale_a_1, + const float& l2_sf) { + const float r0 = (prev_scale_a_0 * prev_l2_sf) * + (kFastMath ? math::fast_rcp(scale_a_0 * l2_sf) : 1.0f / (scale_a_0 * l2_sf)); + const float r1 = (prev_scale_a_1 * prev_l2_sf) * + (kFastMath ? math::fast_rcp(scale_a_1 * l2_sf) : 1.0f / (scale_a_1 * l2_sf)); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] *= r0; + final_accum[i*4+1] *= r0; + final_accum[i*4+2] *= r1; + final_accum[i*4+3] *= r1; + } + }; + auto rescale_l2_act_final = [&](const float& prev_scale_a_0, const float& prev_scale_a_1, + const float& scale_a_0, const float& scale_a_1) { + const float r0 = prev_scale_a_0 * (kFastMath ? math::fast_rcp(scale_a_0) : 1.0f / scale_a_0); + const float r1 = prev_scale_a_1 * (kFastMath ? math::fast_rcp(scale_a_1) : 1.0f / scale_a_1); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] *= r0; + final_accum[i*4+1] *= r0; + final_accum[i*4+2] *= r1; + final_accum[i*4+3] *= r1; + } + }; + + if constexpr (kHidden >= 7168) { + float prev_scale_a_0 = 1.0f, prev_scale_a_1 = 1.0f; + float prev_gate_sf = 1.0f, prev_up_sf = 1.0f, prev_l2_sf = 1.0f; + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + + float scale_a_0_lo, scale_a_1_lo; + float scale_a_0_hi, scale_a_1_hi; + if (block_phase == sched::BlockPhase::Linear1) { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + } else { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); + scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); + scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); + } + + constexpr uint32_t kL1SFKBlocks = kHidden / 128; + constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; + constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; + float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; + if (block_phase == sched::BlockPhase::Linear1) { + const uint32_t gate_n = sf_n_block_idx / 2u; + const uint32_t up_n = kL1SFGateBlks + gate_n; + const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; + gate_sf = __ldg(base + gate_n * kL1SFKBlocks); + up_sf = __ldg(base + up_n * kL1SFKBlocks); + } else { + l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert + + sf_n_block_idx * kL2SFKBlocks + k_block_idx); + } + + if (block_phase == sched::BlockPhase::Linear1) { + if (k_block_idx != 0) + rescale_l1_final(prev_scale_a_0, prev_scale_a_1, + prev_gate_sf, prev_up_sf, + scale_a_0_lo, scale_a_1_lo, + gate_sf, up_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + prev_scale_a_0 = scale_a_0_lo; + prev_scale_a_1 = scale_a_1_lo; + prev_gate_sf = gate_sf; + prev_up_sf = up_sf; + } else { + if (k_block_idx != 0) + rescale_l2_final(prev_scale_a_0, prev_scale_a_1, prev_l2_sf, + scale_a_0_lo, scale_a_1_lo, l2_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + rescale_l2_act_final(scale_a_0_lo, scale_a_1_lo, + scale_a_0_hi, scale_a_1_hi); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + const uint32_t k_off = (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + prev_scale_a_0 = scale_a_0_hi; + prev_scale_a_1 = scale_a_1_hi; + prev_l2_sf = l2_sf; + } + } + + if (num_k_blocks != 0) { + if (block_phase == sched::BlockPhase::Linear1) { + postscale_l1_final(prev_scale_a_0, prev_scale_a_1, + prev_gate_sf, prev_up_sf); + } else { + postscale_l2_final(prev_scale_a_0, prev_scale_a_1, prev_l2_sf); + } + } + } else { + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + + float scale_a_0_lo, scale_a_1_lo; + float scale_a_0_hi, scale_a_1_hi; + if (block_phase == sched::BlockPhase::Linear1) { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + } else { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); + scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); + scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); + } + + constexpr uint32_t kL1SFKBlocks = kHidden / 128; + constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; + constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; + float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; + if (block_phase == sched::BlockPhase::Linear1) { + const uint32_t gate_n = sf_n_block_idx / 2u; + const uint32_t up_n = kL1SFGateBlks + gate_n; + const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; + gate_sf = __ldg(base + gate_n * kL1SFKBlocks); + up_sf = __ldg(base + up_n * kL1SFKBlocks); + } else { + l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert + + sf_n_block_idx * kL2SFKBlocks + k_block_idx); + } + + if (block_phase == sched::BlockPhase::Linear1) { + if (k_block_idx != 0) + prescale_l1_final(scale_a_0_lo, scale_a_1_lo, gate_sf, up_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + postscale_l1_final(scale_a_0_lo, scale_a_1_lo, gate_sf, up_sf); + } else { + if (k_block_idx != 0) + prescale_l2_final(scale_a_0_lo, scale_a_1_lo, l2_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + postscale_l2_final(scale_a_0_lo, scale_a_1_lo, l2_sf); + prescale_l2_final(scale_a_0_hi, scale_a_1_hi, l2_sf); + + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + const uint32_t k_off = (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, final_accum, true); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(final_accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + postscale_l2_final(scale_a_0_hi, scale_a_1_hi, l2_sf); + } + } + } + } else { + float accum[kAccumPerThread]; + + for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + full_barriers[stage_idx]->wait(phase); + + // Read SF (must precede warpgroup_arrive) + float scale_a_0_lo, scale_a_1_lo; + float scale_a_0_hi, scale_a_1_hi; // Only used in L2 (per-64 K) + if (block_phase == sched::BlockPhase::Linear1) { + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + row_offset_r1); + } else { + // L2: SFA layout is (K=2, M=BLOCK_M) MN-major; first half SF at offset 0, second at BLOCK_M + scale_a_0_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r0); + scale_a_1_lo = ptx::ld_shared(smem_sfa[stage_idx] + 0 * BLOCK_M + row_offset_r1); + scale_a_0_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r0); + scale_a_1_hi = ptx::ld_shared(smem_sfa[stage_idx] + 1 * BLOCK_M + row_offset_r1); + } + + // ----- Block (128, 128) weight SF (loaded directly from global) ----- + // L1 weight SF shape: (E, 2*IH/128, H/128) MN-major. The N axis is + // [gate(IH/128), up(IH/128)]; with the gate/up gran-8 interleave on + // the FP8 weight, each logical 128-wide N tile covers 64 rows of gate + // plus 64 rows of up taken from the same original 128-row block, so: + // gate_sf_n = sf_n_block_idx / 2 + // up_sf_n = (IH/128) + sf_n_block_idx / 2 + // + // L2 weight SF shape: (E, H/128, IH/128) MN-major. One scalar per + // logical 128x128 weight-SF tile, broadcast across the matching + // WGMMA accumulators. + // + // Load the weight scale after the barrier from all WG threads. + // This keeps scale loads close to their WGMMA use and lets the + // read-only cache coalesce the same-address accesses. + constexpr uint32_t kL1SFKBlocks = kHidden / 128; + constexpr uint32_t kL2SFKBlocks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFGateBlks = kIntermediateHidden / 128; + constexpr uint32_t kL1SFPerExpert = (kIntermediateHidden * 2 / 128) * kL1SFKBlocks; + constexpr uint32_t kL2SFPerExpert = (kHidden / 128) * kL2SFKBlocks; + float gate_sf = 0.0f, up_sf = 0.0f, l2_sf = 0.0f; + if (block_phase == sched::BlockPhase::Linear1) { + const uint32_t gate_n = sf_n_block_idx / 2u; + const uint32_t up_n = kL1SFGateBlks + gate_n; + const float* base = l1_weights_sf + local_expert_idx * kL1SFPerExpert + k_block_idx; + gate_sf = __ldg(base + gate_n * kL1SFKBlocks); + up_sf = __ldg(base + up_n * kL1SFKBlocks); + } else { + l2_sf = __ldg(l2_weights_sf + local_expert_idx * kL2SFPerExpert + + sf_n_block_idx * kL2SFKBlocks + k_block_idx); + } + + if (block_phase == sched::BlockPhase::Linear1) { + // Single per-128 K-block WGMMA group + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + // L1: gate/up alternate at gran=8 along N; each `i` block of 8 + // cols belongs entirely to one of {gate, up}, so .x and .y + // share the same scalar. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + const float sb = (i & 1u) ? up_sf : gate_sf; + final_accum[i*4+0] += scale_a_0_lo * sb * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_lo * sb * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_lo * sb * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_lo * sb * accum[i*4+3]; + } + } else { + // L2: split BLOCK_K=128 into two halves (per-64 SFA), each 2 WGMMAs. + // First half: K=0..63, SFA = scale_a_*_lo + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k * WGMMA::K, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k * WGMMA::K, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + // L2 first half: single scalar `l2_sf` broadcast across N. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] += scale_a_0_lo * l2_sf * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_lo * l2_sf * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_lo * l2_sf * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_lo * l2_sf * accum[i*4+3]; + } + + // Second half: K=64..127, SFA = scale_a_*_hi + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_arrive(); + #pragma unroll + for (uint32_t k = 0; k < (BLOCK_K / 2) / WGMMA::K; ++ k) { + const uint32_t k_off = (BLOCK_K / 2) + k * WGMMA::K; + auto desc_a = mma::sm90::make_smem_desc( + smem_a[stage_idx] + smem_a_wg_offset + k_off, 1); + auto desc_b = mma::sm90::make_smem_desc( + smem_b[stage_idx] + smem_b_wg_offset + k_off, 1); + WGMMA::wgmma(desc_a, desc_b, accum, k); + } + ptx::warpgroup_commit_batch(); + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread; ++ i) ptx::warpgroup_fence_operand(accum[i]); + ptx::warpgroup_wait<0>(); + + if (lane_idx == 0) + empty_barriers[stage_idx]->arrive(); + + // L2 second half: same broadcast scalar `l2_sf`. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 4; ++ i) { + final_accum[i*4+0] += scale_a_0_hi * l2_sf * accum[i*4+0]; + final_accum[i*4+1] += scale_a_0_hi * l2_sf * accum[i*4+1]; + final_accum[i*4+2] += scale_a_1_hi * l2_sf * accum[i*4+2]; + final_accum[i*4+3] += scale_a_1_hi * l2_sf * accum[i*4+3]; + } + } + } + } + + // Skip epilogue when block is past valid M (still must release via empty) + if (row_base >= valid_m) { + if (block_phase == sched::BlockPhase::Linear1) { + if constexpr (not kL2ArrivalCounter) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } else { + if constexpr (kL2EpilogueRequiresFullSync) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + return; + } + + if (block_phase == sched::BlockPhase::Linear1) { + + // ---------------- L1 EPILOGUE: activation + FP8 quantize + TMA store ---------------- + // Layout in `final_accum`: + // kAccumPerThread/4 chunks, each chunk = 4 floats per thread = + // (r0c0, r0c1, r1c0, r1c1). + // Gate and up chunks alternate; pair `p` uses chunks 2p and 2p+1. + // + // For each pair we produce 4 post-SwiGLU floats per thread, mapped to + // output cols (p*8 + col_idx*2 + {0,1}) for both r0 and r1. + + constexpr uint32_t kNumPairs = kAccumPerThread / 8; + float sf_r0, sf_inv_r0; + float sf_r1, sf_inv_r1; + + float swiglu_r0[kNumPairs][2]; + float swiglu_r1[kNumPairs][2]; + float amax_r0 = 0.0f, amax_r1 = 0.0f; + + auto clamp_gate = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(x, kActivationClamp); + }; + auto clamp_up = [](float& x) { + if constexpr (kActivationClamp != cute::numeric_limits::infinity()) + x = cute::min(cute::max(x, -kActivationClamp), kActivationClamp); + }; + auto silu = [](float x) -> float { + const float e = kFastMath ? __expf(-x) : expf(-x); + const float sig = kFastMath ? math::fast_rcp(1.0f + e) : 1.0f / (1.0f + e); + return x * sig; + }; + + #pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++ p) { + const uint32_t gate = 2 * p, up = 2 * p + 1; + + float g_r0_c0 = final_accum[gate*4 + 0]; + float g_r0_c1 = final_accum[gate*4 + 1]; + float g_r1_c0 = final_accum[gate*4 + 2]; + float g_r1_c1 = final_accum[gate*4 + 3]; + float u_r0_c0 = final_accum[up*4 + 0]; + float u_r0_c1 = final_accum[up*4 + 1]; + float u_r1_c0 = final_accum[up*4 + 2]; + float u_r1_c1 = final_accum[up*4 + 3]; + clamp_gate(g_r0_c0); + clamp_gate(g_r0_c1); + clamp_gate(g_r1_c0); + clamp_gate(g_r1_c1); + clamp_up(u_r0_c0); + clamp_up(u_r0_c1); + clamp_up(u_r1_c0); + clamp_up(u_r1_c1); + + if (valid_r0) { + swiglu_r0[p][0] = silu(g_r0_c0) * u_r0_c0; + swiglu_r0[p][1] = silu(g_r0_c1) * u_r0_c1; + amax_r0 = cute::max(amax_r0, cute::max(cute::abs(swiglu_r0[p][0]), cute::abs(swiglu_r0[p][1]))); + } else { + swiglu_r0[p][0] = 0.0f; + swiglu_r0[p][1] = 0.0f; + } + if (valid_r1) { + swiglu_r1[p][0] = silu(g_r1_c0) * u_r1_c0; + swiglu_r1[p][1] = silu(g_r1_c1) * u_r1_c1; + amax_r1 = cute::max(amax_r1, cute::max(cute::abs(swiglu_r1[p][0]), cute::abs(swiglu_r1[p][1]))); + } else { + swiglu_r1[p][0] = 0.0f; + swiglu_r1[p][1] = 0.0f; + } + } + + // Apply token weight: SwiGLU * topk_weight (single load per row) + const float weight_r0 = valid_r0 ? *l1_topk_weights_buffer + .get_data_buffer(m_idx + row_offset_r0) + .get_base_ptr() : 0.0f; + const float weight_r1 = valid_r1 ? *l1_topk_weights_buffer + .get_data_buffer(m_idx + row_offset_r1) + .get_base_ptr() : 0.0f; + #pragma unroll + for (uint32_t p = 0; p < kNumPairs; ++ p) { + swiglu_r0[p][0] *= weight_r0; + swiglu_r0[p][1] *= weight_r0; + swiglu_r1[p][0] *= weight_r1; + swiglu_r1[p][1] *= weight_r1; + } + + amax_r0 *= cute::abs(weight_r0); + amax_r1 *= cute::abs(weight_r1); + + // Reduce amax across the 4 col-lanes that share the same row. In the + // SM90 WGMMA output layout, lanes with the same `lane_idx >> 2` and + // different `lane_idx & 3` partition the WG-owned output columns for + // the same r_0/r_1, so this is an INTRA-group reduction + // (`warp_reduce<4, false>`). Using `<4, true>` would instead merge + // amax across 8 different rows -- giving wrong per-row SF. + amax_r0 = math::warp_reduce<4, false>(amax_r0, math::ReduceMax()); + amax_r1 = math::warp_reduce<4, false>(amax_r1, math::ReduceMax()); + + // Phase 2: cross-WG amax. When two N-split warpgroups share one + // per-64 SF group, each WG so far only saw its own + // WG_L1_OUT_BLOCK_N columns; the true per-row amax spans both + // halves. Reduce across both warpgroups through a small smem + // scratch (carved from the upper, currently-unused half of the + // CD staging region) so BOTH WGs quantize with the SAME SF. + if constexpr (kSplitNSharesSF) { + float* amax_scratch = reinterpret_cast( + reinterpret_cast(smem_cd_l1) + SMEM_CD_SIZE / 2); + #pragma unroll + for (uint32_t i = epilogue_thread_idx; i < BLOCK_M; i += kNumEpilogueThreads) + amax_scratch[i] = 0.0f; + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (col_idx == 0) { + atomicMax(reinterpret_cast(&amax_scratch[r_0]), __float_as_uint(amax_r0)); + atomicMax(reinterpret_cast(&amax_scratch[r_1]), __float_as_uint(amax_r1)); + } + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + amax_r0 = amax_scratch[r_0]; + amax_r1 = amax_scratch[r_1]; + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + + // Compute SF and inverse SF for each row + float2 amax_pair = {amax_r0, amax_r1}; + float2 sf_pair, sf_inv_pair; + sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv(amax_pair, sf_pair, sf_inv_pair); + sf_r0 = sf_pair.x; sf_inv_r0 = sf_inv_pair.x; + sf_r1 = sf_pair.y; sf_inv_r1 = sf_inv_pair.y; + + // Quantize and write to the shared-memory staging tile. + auto* smem_cd_l1_wg = smem_cd_l1 + smem_cd_l1_wg_offset; + DG_STATIC_ASSERT(kNumPairs % 2 == 0, "L1 staging stores two 8-byte chunks at once"); + #pragma unroll + for (uint32_t p_base = 0; p_base < kNumPairs; p_base += 2) { + uint16_t r0_bits[2], r1_bits[2]; + #pragma unroll + for (uint32_t q = 0; q < 2; ++ q) { + const uint32_t p = p_base + q; + const float v00 = swiglu_r0[p][0] * sf_inv_r0; + const float v01 = swiglu_r0[p][1] * sf_inv_r0; + const float v10 = swiglu_r1[p][0] * sf_inv_r1; + const float v11 = swiglu_r1[p][1] * sf_inv_r1; + + const __nv_fp8x2_e4m3 r0_pair(make_float2(v00, v01)); + const __nv_fp8x2_e4m3 r1_pair(make_float2(v10, v11)); + r0_bits[q] = valid_r0 ? r0_pair.__x : 0u; + r1_bits[q] = valid_r1 ? r1_pair.__x : 0u; + } + + #pragma unroll + for (uint32_t q = 0; q < 2; ++ q) { + const uint32_t p = p_base + q; + const uint32_t col = p * 8 + col_idx * 2; + auto* p0 = reinterpret_cast( + smem_cd_l1_wg + r_0 * WG_SMEM_CD_L1_STRIDE_N + col); + auto* p1 = reinterpret_cast( + smem_cd_l1_wg + r_1 * WG_SMEM_CD_L1_STRIDE_N + col); + if (valid_r0) + *p0 = r0_bits[q]; + if (valid_r1) + *p1 = r1_bits[q]; + } + } + + // Write SF as float at `[token, n_block_idx]` in L2 acts SF buffer (per-64 layout). + // Each row is contributed by lanes col_idx in {0..3}; only col_idx == 0 writes. + // In the shared-SF split both warpgroups own the same per-64 group and rows, so + // only the first N-split warpgroup publishes the SF slot to avoid a write race. + if (col_idx == 0 and (not kSplitNSharesSF or epilogue_wg_n_idx == 0)) { + auto sf_base_ptr = l2_sf_buffer.get_base_ptr(); + // SF buffer is (kNumPaddedSFPoolTokens x kIntermediateHidden/64), MN-major: + // addr[k_idx * num_padded_sf_pool_tokens + token_idx] + const uint32_t token_r0 = pool_block_idx * BLOCK_M + row_offset_r0; + const uint32_t token_r1 = pool_block_idx * BLOCK_M + row_offset_r1; + const uint32_t k_sf_idx = sf_n_block_idx; // one per-64 post-SwiGLU group + if (valid_r0) + sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r0] = sf_r0; + if (valid_r1) + sf_base_ptr[k_sf_idx * kNumPaddedSFPoolTokens + token_r1] = sf_r1; + } + + // Sync the warpgroup before TMA store. In the shared-tile split + // both N-split warpgroups must finish writing their halves of the + // joint L1-output tile, so sync across all epilogue threads. + if constexpr (kSplitNSharesSF) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + else + ptx::sync_aligned(128, kEpilogueWGBarrierStartIdx + epilogue_wg_idx); + + // Issue TMA store of the entire tile. Padding rows beyond + // `valid_m` are written with stale/garbage FP8 to the L1-output + // pool buffer, but they are never consumed downstream: the L2 + // GEMM tile loads them, but its NVLink-scatter epilogue is + // gated by `m_idx_in_block >= valid_m`, and stale SF in the + // padding rows can produce NaN accumulators that simply stay + // in registers (only valid rows are converted to BF16 and + // STSM'd into smem). Using TMA for partial tiles is a large + // win for low-batch / decode where every tile is partial. + if constexpr (kSplitNSharesSF) { + // One combined store of the joint L1_OUT_BLOCK_N tile, issued + // by the first N-split warpgroup once both halves are staged. + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { + const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy( + &tensor_map_l1_output, + smem_cd_l1, + out_n_idx, + m_idx + row_base); + cute::tma_store_arrive(); + } + } else { + if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + const uint32_t out_n_idx = n_block_idx * L1_OUT_BLOCK_N + wg_l1_out_n_offset; + cute::tma_store_fence(); + cute::SM90_TMA_STORE_2D::copy( + &tensor_map_l1_output, + smem_cd_l1 + smem_cd_l1_wg_offset, + out_n_idx, + m_idx + row_base); + cute::tma_store_arrive(); + } + } + __syncwarp(); + ptx::tma_store_wait<0>(); + + // Notify L2 that this L1 output (and SF) is ready. Counter mode lets + // independent WG tiles publish arrivals without the CTA-wide barrier + // needed before the single bit-mask update. + if constexpr (kL2ArrivalCounter) { + if constexpr (kSplitNSharesSF) { + // The combined tile counts for both N-split warpgroups; the + // storing warpgroup publishes all kWarpgroupSplitN arrivals + // after its TMA store has drained. + if (epilogue_wg_n_idx == 0 and warp_idx_in_wg == 0 and cute::elect_one_sync()) { + ptx::red_add_rel( + reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), + kWarpgroupSplitN); + } + } else if (warp_idx_in_wg == 0 and cute::elect_one_sync()) { + ptx::red_add_rel( + reinterpret_cast(workspace.get_l2_arrival_mask_ptr(pool_block_idx)), 1); + } + } else { + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { + ptx::red_or_rel_gpu( + workspace.get_l2_arrival_mask_ptr(pool_block_idx), + 1ull << n_block_idx); + } + } + __syncwarp(); + // In the shared-tile split only the first warpgroup issues and + // drains the combined TMA store; gate the other warpgroup so it + // cannot overwrite the joint smem tile in the next block until + // that store has drained. + if constexpr (kSplitNSharesSF) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } else { + // ---------------- L2 EPILOGUE: BF16 cast + NVLink scatter ---------------- + constexpr uint32_t kNumRowsPerWarp = WG_BLOCK_M / 8; + + const uint32_t row_in_warp_block = lane_idx / 16; // 0 or 1 + const uint32_t lane_in_row = lane_idx % 16; + const uint32_t cols_per_lane = WG_BLOCK_N / 16; + + // STSM into smem_cd_l2 (BF16). Reuse SM100 column-swizzle layout. + #pragma unroll + for (uint32_t i = 0; i < kAccumPerThread / 8; ++ i) { + // Each i consumes 8 floats (one 16x256b chunk in SM100 terms). + // For SM90 WGMMA layout, 8 floats per i correspond to 2 chunks of 4 floats: + // final_accum[i*8 + (0..3)] = chunk 2i: (r0c0, r0c1, r1c0, r1c1) + // final_accum[i*8 + (4..7)] = chunk 2i+1: same shape + const uint32_t chunk_lo = 2 * i, chunk_hi = 2 * i + 1; + + auto write_pair = [&](uint32_t row, uint32_t col, uint32_t packed) { + auto smem_ptr = smem_cd_l2 + + smem_cd_l2_wg_offset + + row * WG_BLOCK_N + + col; + // BF16 STS: 2 bf16 elements + *reinterpret_cast(smem_ptr) = packed; + }; + if (valid_r0) { + const uint32_t r0_lo = math::cast_into_bf16_and_pack( + final_accum[chunk_lo*4 + 0], final_accum[chunk_lo*4 + 1]); + const uint32_t r0_hi = math::cast_into_bf16_and_pack( + final_accum[chunk_hi*4 + 0], final_accum[chunk_hi*4 + 1]); + write_pair(r_0, chunk_lo * 8 + col_idx * 2, r0_lo); + write_pair(r_0, chunk_hi * 8 + col_idx * 2, r0_hi); + } + if (valid_r1) { + const uint32_t r1_lo = math::cast_into_bf16_and_pack( + final_accum[chunk_lo*4 + 2], final_accum[chunk_lo*4 + 3]); + const uint32_t r1_hi = math::cast_into_bf16_and_pack( + final_accum[chunk_hi*4 + 2], final_accum[chunk_hi*4 + 3]); + write_pair(r_1, chunk_lo * 8 + col_idx * 2, r1_lo); + write_pair(r_1, chunk_hi * 8 + col_idx * 2, r1_hi); + } + } + + // Each warp writes and then scatters only its own 16-row + // slice, so a warp-level fence is enough before reading + // back from shared memory. + __syncwarp(); + + // Scatter to remote ranks via NVLink (one row per warp-pair) + // Each warpgroup-warp covers 8 unique rows x 2 (r_0 + r_1 doubled by warps) + // Lane group of 16 within a warp -> 1 row. + // Each lane copies `cols_per_lane` BF16 (= cols_per_lane*2 bytes) as one + // vector. WG_BLOCK_N=128 -> 8 BF16 = uint4; WG_BLOCK_N=64 -> 4 BF16 = uint2. + using ScatterVec = std::conditional_t<(WG_BLOCK_N <= 64), uint2, uint4>; + DG_STATIC_ASSERT(cols_per_lane * sizeof(nv_bfloat16) == sizeof(ScatterVec), + "Scatter vector width must match cols_per_lane"); + #pragma unroll + for (uint32_t j = 0; j < kNumRowsPerWarp; ++ j) { + const uint32_t row_in_wg = warp_idx_in_wg * 16 + j * 2 + row_in_warp_block; + const uint32_t m_idx_in_block = row_base + row_in_wg; + if (m_idx_in_block >= valid_m) break; + + // Read cols_per_lane BF16 (= one ScatterVec) from smem + auto smem_ptr = smem_cd_l2 + + smem_cd_l2_wg_offset + + row_in_wg * WG_BLOCK_N + + lane_in_row * cols_per_lane; + const auto packed = *reinterpret_cast(smem_ptr); + + const auto src_metadata = *workspace.get_token_src_metadata_ptr(m_idx + m_idx_in_block); + const uint32_t dst_rank_idx = src_metadata.rank_idx; + const uint32_t dst_token_idx = src_metadata.token_idx; + const uint32_t dst_topk_idx = src_metadata.topk_idx; + const auto dst_token = combine_token_buffer.get_rank_buffer(dst_topk_idx) + .get_data_buffer(dst_token_idx); + auto dst_ptr = math::advance_ptr( + dst_token.get_base_ptr(), + (n_idx + wg_n_offset) * sizeof(nv_bfloat16) + lane_in_row * sizeof(ScatterVec)); + *sym_buffer.map(dst_ptr, dst_rank_idx) = packed; + } + + if constexpr (kL2EpilogueRequiresFullSync) + ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); + } + }; + + if constexpr (kSplitPhaseHotPath) { + sm90_fp8_mega_moe_for_each_block_split( + scheduler, + [&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_math_block( + std::integral_constant{}, + local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }, + [&](const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_math_block( + std::integral_constant{}, + local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }); + } else { + scheduler.for_each_block([&](const sched::BlockPhase& block_phase, + const uint32_t& local_expert_idx, + const uint32_t& num_k_blocks, + const uint32_t& m_block_idx, const uint32_t& n_block_idx) { + process_math_block(block_phase, local_expert_idx, num_k_blocks, m_block_idx, n_block_idx); + }); + } + + // ---------------- COMBINE ---------------- + // NVLink barrier first: signals remote ranks that this rank's GEMM + // outputs (NVLink scatter targets) are fully written. + comm::nvlink_barrier( + workspace, sym_buffer, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + + // Sync with dispatch (paired with dispatch's pre-cleanup sync) so that + // dispatch may now safely clean workspace state. + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + if (epilogue_warp_idx >= kNumCombineWarps) + return; + + constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); + constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); + + constexpr uint32_t kNumChunkSlots = 3; + constexpr uint32_t kNumMaxRegistersForBuffer = 128; + constexpr uint32_t kDefaultNumChunks = + (kNumChunkSlots * kNumCombineWarps * kNumHiddenBytes <= SMEM_BEFORE_BARRIER_SIZE + and kHidden <= 32 * kNumMaxRegistersForBuffer) ? 1 : 2; + // Flash-style hidden=7168 is 7 * 1024. Splitting combine into 7 chunks + // keeps each lane's BF16 reduce accumulator much smaller without + // violating the 32-lane uint4 mapping. + constexpr uint32_t kSplitMNNumChunks = (kHidden % 7 == 0) ? 7 : (kHidden >= 1024 ? 4 : 1); + constexpr uint32_t kNumChunks = kSplitMNWarpgroups ? kSplitMNNumChunks : kDefaultNumChunks; + constexpr uint32_t kNumChunkBytes = kNumHiddenBytes / kNumChunks; + constexpr uint32_t kNumChunkUint4 = kNumChunkBytes / sizeof(uint4); + constexpr uint32_t kNumUint4PerLane = kNumChunkUint4 / 32; + DG_STATIC_ASSERT(kHidden % kNumChunks == 0, "Hidden must be divisible by number of chunks"); + DG_STATIC_ASSERT(kNumChunkSlots * kNumCombineWarps * kNumHiddenBytes / kNumChunks <= SMEM_BEFORE_BARRIER_SIZE, "Hidden is too large"); + DG_STATIC_ASSERT(kNumChunkBytes % 16 == 0, "Combine chunk must be TMA-aligned (16 bytes)"); + DG_STATIC_ASSERT(kNumChunkBytes % sizeof(uint4) == 0, "Combine chunk must be divisible by 16 bytes"); + DG_STATIC_ASSERT(kNumChunkUint4 % 32 == 0, "Combine chunk must be a multiple of 32 16-byte elements"); + DG_STATIC_ASSERT(kNumTopk <= 32, "Top-k must fit in a single warp"); + + DG_TRAP_ONLY_DEVICE_ASSERT(kNumChunkSlots * kNumCombineWarps * kNumChunkBytes <= static_cast( + reinterpret_cast(barrier_start_ptr) - smem_buffer)); + + const auto combine_load_buffer = utils::PatternVisitor([&](const uint32_t& i) { + return math::advance_ptr(smem_buffer, (epilogue_warp_idx + i * kNumCombineWarps) * kNumChunkBytes); + }); + const auto combine_store_buffer = math::advance_ptr( + smem_buffer, (epilogue_warp_idx + kNumCombineWarps * 2) * kNumChunkBytes); + + auto combine_load_barriers = utils::PatternVisitor([&](const uint32_t& i) { + return combine_barriers[i + epilogue_warp_idx * 2]; + }); + + uint32_t combine_phase = 0; + uint32_t load_stage_idx = 0; + for (uint32_t token_idx = sm_idx * kNumCombineWarps + epilogue_warp_idx; + token_idx < num_tokens; + token_idx += kNumSMs * kNumCombineWarps) { + const int stored_topk_slot_idx = lane_idx < kNumTopk ? + static_cast(__ldg(input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : -1; + const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); + + for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { + const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; + + uint32_t mask = total_mask; + const auto move_mask_and_load = [&](const uint32_t& i) { + if (mask) { + const uint32_t slot_idx = __ffs(mask) - 1; + mask ^= 1 << slot_idx; + if (cute::elect_one_sync()) { + const auto src_ptr = math::advance_ptr( + combine_token_buffer.get_rank_buffer(slot_idx) + .get_data_buffer(token_idx).get_base_ptr(), + chunk_byte_offset); + ptx::tma_load_1d(combine_load_buffer[i], src_ptr, combine_load_barriers[i], kNumChunkBytes); + ptx::mbarrier_arrive_and_set_tx(combine_load_barriers[i], kNumChunkBytes); + } + __syncwarp(); + return true; + } + return false; + }; + + bool do_reduce = move_mask_and_load(load_stage_idx); + + float2 reduced[kNumUint4PerLane * kNumElemsPerUint4] = {}; + while (do_reduce) { + do_reduce = move_mask_and_load(load_stage_idx ^ 1); + combine_load_barriers[load_stage_idx]->wait(combine_phase); + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + const auto uint4_values = combine_load_buffer[load_stage_idx][j * 32 + lane_idx]; + const auto bf16_values = reinterpret_cast(&uint4_values); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + ptx::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); + } + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } + + #pragma unroll + for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { + uint4 casted; + auto casted_bf16 = reinterpret_cast(&casted); + #pragma unroll + for (uint32_t l = 0; l < kNumElemsPerUint4; ++ l) + casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); + + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); + } + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, + casted.x, casted.y, casted.z, casted.w); + } + __syncwarp(); + + if (cute::elect_one_sync()) { + cute::tma_store_fence(); + ptx::tma_store_1d( + math::advance_ptr(y, static_cast(token_idx) * kNumHiddenBytes + chunk_byte_offset), + combine_store_buffer, kNumChunkBytes); + cute::tma_store_arrive(); + } + __syncwarp(); + } + } + } +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_90"); +#endif +} + +} // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index 12cbf9ffea..21bf6a7ed1 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -3,7 +3,7 @@ import os import shutil import subprocess -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Optional, Tuple import torch import tvm_ffi @@ -270,14 +270,149 @@ def k_grouped_bf16_gemm_tn_contiguous(a, b, d, ks, ks_tensor, c=None, compiled_d pass # Mega kernels +from . import mega from .mega import ( SymmBuffer, - get_symm_buffer_for_mega_moe, transform_weights_for_mega_moe, fp8_fp4_mega_moe, mega_moe_pre_dispatch, ) + +def _from_dlpack_if_needed(tensor, dtype: Optional[torch.dtype] = None) -> torch.Tensor: + if not isinstance(tensor, torch.Tensor): + tensor = torch.utils.dlpack.from_dlpack(tensor) + if dtype is not None and tensor.dtype != dtype: + tensor = tensor.view(dtype) + return tensor + + +class SM90SymmBuffer: + def __init__(self, group, + num_experts: int, + num_max_tokens_per_rank: int, num_topk: int, + hidden: int, intermediate_hidden: int, + use_fp8_dispatch: bool = True, + activation: str = 'swiglu'): + import torch.distributed._symmetric_memory as symm_mem + + self.group = group + self.num_experts = num_experts + self.num_max_tokens_per_rank = num_max_tokens_per_rank + self.num_topk = num_topk + self.hidden = hidden + self.intermediate_hidden = intermediate_hidden + + num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_sm90_mega_moe( + group.size(), num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + use_fp8_dispatch, activation + ) + self.buffer = symm_mem.empty(num_bytes, dtype=torch.int8, device='cuda') + self.handle = symm_mem.rendezvous(self.buffer, group=group) + self.buffer.zero_() + self.group.barrier() + torch.cuda.synchronize() + + (x, x_sf, topk_idx, topk_weights, + l1_acts, l1_acts_sf, l2_acts, l2_acts_sf) = slice_input_buffers(self.buffer) + self.x = _from_dlpack_if_needed(x, torch.float8_e4m3fn) + self.x_sf = _from_dlpack_if_needed(x_sf) + self.topk_idx = _from_dlpack_if_needed(topk_idx) + self.topk_weights = _from_dlpack_if_needed(topk_weights) + self.l1_acts = _from_dlpack_if_needed(l1_acts, torch.float8_e4m3fn) + self.l1_acts_sf = _from_dlpack_if_needed(l1_acts_sf) + self.l2_acts = _from_dlpack_if_needed(l2_acts, torch.float8_e4m3fn) + self.l2_acts_sf = _from_dlpack_if_needed(l2_acts_sf) + + def destroy(self): + self.handle = None + self.buffer = None + self.group = None + self.x = None + self.x_sf = None + + +def get_symm_buffer_for_sm90_mega_moe(group, + num_experts: int, + num_max_tokens_per_rank: int, num_topk: int, + hidden: int, intermediate_hidden: int, + use_fp8_dispatch: bool = True, + activation: str = 'swiglu') -> SM90SymmBuffer: + from .utils.math import align + + num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe()) + return SM90SymmBuffer( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + use_fp8_dispatch, activation + ) + + +def get_symm_buffer_for_mega_moe(group, + num_experts: int, + num_max_tokens_per_rank: int, num_topk: int, + hidden: int, intermediate_hidden: int, + use_fp8_dispatch: bool = True, + activation: str = 'swiglu'): + if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9: + return get_symm_buffer_for_sm90_mega_moe( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + use_fp8_dispatch, activation + ) + return mega.get_symm_buffer_for_mega_moe( + group, num_experts, + num_max_tokens_per_rank, num_topk, + hidden, intermediate_hidden, + use_fp8_dispatch, activation + ) + + +def transform_weights_for_mega_moe_sm90( + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor] +) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + l1_fp8, l1_sf = l1_weights + + def _interleave_one(t, gran: int = 8) -> torch.Tensor: + g, n, *rest = t.shape + half = n // 2 + gate = t[:, :half].reshape(g, half // gran, gran, *rest) + up = t[:, half:].reshape(g, half // gran, gran, *rest) + return torch.empty_like(t).copy_(torch.stack([gate, up], dim=2).reshape(g, n, *rest)) + + return (_interleave_one(l1_fp8), l1_sf), l2_weights + + +def fp8_mega_moe(y: torch.Tensor, + l1_weights: Tuple[torch.Tensor, torch.Tensor], + l2_weights: Tuple[torch.Tensor, torch.Tensor], + sym_buffer: SM90SymmBuffer, + cumulative_local_expert_recv_stats: Optional[torch.Tensor] = None, + recipe: Tuple[int, int, int] = (128, 128, 128), + activation: str = 'swiglu', + activation_clamp: Optional[float] = None, + fast_math: bool = True): + (l1_weights_data, l1_weights_sf) = l1_weights + (l2_weights_data, l2_weights_sf) = l2_weights + _C.fp8_mega_moe( + y, + l1_weights_data, l1_weights_sf, + l2_weights_data, l2_weights_sf, + cumulative_local_expert_recv_stats, + sym_buffer.buffer, + sym_buffer.handle.buffer_ptrs, sym_buffer.group.rank(), + sym_buffer.num_max_tokens_per_rank, + sym_buffer.num_experts, sym_buffer.num_topk, + recipe, + activation, activation_clamp, + fast_math + ) + # Some utils from . import testing from . import utils diff --git a/tests/test_mega_moe_hopper.py b/tests/test_mega_moe_hopper.py new file mode 100644 index 0000000000..f3579a7e91 --- /dev/null +++ b/tests/test_mega_moe_hopper.py @@ -0,0 +1,1946 @@ +"""SM90 (Hopper) MegaMoE fused-kernel and same-pipeline baseline benchmark. + +This follows the structure of ``tests/test_mega_moe.py`` for the SM100 FP4 +path, with the compute path changed to SM90 FP8: + +* fused: calls ``deep_gemm.fp8_mega_moe`` (kernel symbol + ``sm90_fp8_mega_moe_impl``) with weights transformed by + ``transform_weights_for_mega_moe_sm90`` and a ``SymmBuffer``. +* baseline: DeepEP dispatch, two grouped FP8 GEMMs, Triton SwiGLU, and DeepEP + combine with untransformed weights. The current SM90 grouped GEMM path accepts + per-128-K L2 activation SF, while the fused SM90 MegaMoE L1 epilogue writes + per-64-K L2 activation SF to avoid cross-CTA synchronization. This is a + same-pipeline performance reference, not a bitwise correctness oracle. +* low-latency baseline (optional, ``--run-low-latency-baseline``): mirrors the + sglang low-latency MoE pipeline (see + ``sglang/srt/layers/moe/token_dispatcher/deepep.py::_DeepEPDispatcherImplLowLatency``): + ``Buffer.low_latency_dispatch`` (use_fp8=True) -> per-expert masked-layout + FP8 grouped GEMM -> masked SwiGLU + FP8 quant -> masked FP8 grouped GEMM -> + ``Buffer.low_latency_combine`` (which applies topk weights internally). This + is the canonical decode path used in production EP serving. +* fused-only sweep (optional, ``--fused-only-sweep``): replaces the old + standalone SM90 benchmark harness. It sweeps token counts, measures + only the fused SM90 kernel, and keeps the ``--ncu-profile-only`` / + ``--local-rank-idx`` interface for single-rank NCU profiling. +* accuracy mode (optional, ``--accuracy``): runs the former layered SM90 + correctness suite with a PyTorch BF16/FP32 reference. It covers smoke, + heuristic branches, shape sweeps, edge cases, and optional random stress. +* output: TFLOPS, overlap-adjusted TFLOPS, HBM GB/s, NVLink GB/s, fused time, + reduction estimate, and ``t_baseline / t_fused``. +""" + +import argparse +import math +import os +import random +import sys +import torch +import torch.distributed as dist +import triton +import triton.language as tl +from typing import Tuple, List, Dict, Any + +REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +if os.getenv("DG_TEST_USE_SOURCE_TREE", "0") == "1" and REPO_ROOT not in sys.path: + sys.path.insert(0, REPO_ROOT) + +import deep_gemm +from deep_gemm.utils import per_token_cast_to_fp8 +from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather +from deep_gemm.testing import bench_kineto, calc_diff, get_arch_major + +_missing_sgl_symbols = [ + name for name in ("fp8_mega_moe", "transform_weights_for_mega_moe_sm90") + if not hasattr(deep_gemm, name) +] +if _missing_sgl_symbols: + raise RuntimeError( + "SM90 MegaMoE tests require the sgl-deep-gemm wheel built with " + "`bash build_sgl_deep_gemm.sh`; missing symbols: " + f"{', '.join(_missing_sgl_symbols)}" + ) + +try: + import deep_ep as _deep_ep + _deep_ep_import_error = None +except Exception as ex: + _deep_ep = None + _deep_ep_import_error = ex + + +# Must match the template entry point in +# deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh so bench_kineto can +# select the fused MegaMoE GPU region from the trace. +SM90_KERNEL_NAME = "sm90_fp8_mega_moe_impl" + + +# Max finite value of FP8 e4m3fn; quantization uses amax / 448 as the scale. +FP8_E4M3_MAX = 448.0 +# Triton >= 3 requires Python globals read by a JIT kernel to be tl.constexpr, +# otherwise compilation can fail with NameError. Host-side torch code still uses +# the plain float above. +_FP8_E4M3_MAX_TL = tl.constexpr(448.0) +L1_ACT_SF_GRAN = 128 +FUSED_L2_ACT_SF_GRAN = 64 +BASELINE_L2_ACT_SF_GRAN = 128 +WEIGHT_SF_GRAN_MN = 128 +WEIGHT_SF_GRAN_K = 128 + + +# ============================================================================ +# Section 1: Triton SwiGLU + FP8 quantization kernel. +# ---------------------------------------------------------------------------- +# The baseline L2 path uses DeepGEMM SM90 grouped FP8 GEMM, which accepts +# per-128-K activation SF. The scale values still use the same power-of-two +# rounding as the fused epilogue to avoid adding an exact-FP32-scale difference. +# Input x : (M, 2*H) bf16, laid out as [gate_part | up_part]. +# Input topk_w : (M,) fp32, optional. +# Output y : (M, H) fp8_e4m3fn. +# Output y_sf : (M, H / BLOCK_K) fp32, row-major. +# ============================================================================ + + +@triton.jit +def _swiglu_apply_weight_to_fp8_kernel( + x_ptr, + topk_w_ptr, + y_ptr, + y_sf_ptr, + M, + H, # Runtime shape + stride_xm, + stride_xn, # x: (M, 2H) stride + stride_ym, + stride_yn, # y: (M, H) stride + stride_sfm, + stride_sfk, # y_sf: (M, H / BLOCK_K) stride + clamp_value, # Ignored when HAS_CLAMP=False + HAS_TOPK: tl.constexpr, + HAS_CLAMP: tl.constexpr, + USE_UE8M0_SCALE: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_K: tl.constexpr, # = num_per_channels +): + # One program handles BLOCK_M tokens and one BLOCK_K column tile. + pid_m = tl.program_id(0) + pid_k = tl.program_id(1) + + # Row indices handled by this program. + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + # Column indices inside the current K block, in the H dimension. + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + mask_m = offs_m < M + + # 1) Load gate from [0, H) and up from [H, 2H). + # stride_xn is an element stride, so H + offs_k is also element-based. + gate_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xn + up_ptrs = x_ptr + offs_m[:, None] * stride_xm + (H + offs_k[None, :]) * stride_xn + gate = tl.load(gate_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + up = tl.load(up_ptrs, mask=mask_m[:, None], other=0.0).to(tl.float32) + + # 2) Optional clamp: one-sided for gate, two-sided for up. + if HAS_CLAMP: + gate = tl.minimum(gate, clamp_value) + up = tl.minimum(tl.maximum(up, -clamp_value), clamp_value) + + # 3) SwiGLU: silu(gate) * up = gate * sigmoid(gate) * up, accumulated in FP32. + y = gate * tl.sigmoid(gate) * up + + # 4) Optional MoE weight scaling with a per-token scalar. + if HAS_TOPK: + w = tl.load(topk_w_ptr + offs_m, mask=mask_m, other=1.0) + y = y * w[:, None] + + # 5) Per-row absmax in the current K block -> scale. + amax = tl.max(tl.abs(y), axis=1) # (BLOCK_M,) + sf = tl.maximum(amax / _FP8_E4M3_MAX_TL, 1.0e-30) + if USE_UE8M0_SCALE: + # Match deep_gemm/common/math.cuh::get_e4m3_sf_and_sf_inv: + # scale = 2 ** ceil(log2(amax / 448)). + sf = tl.exp2(tl.ceil(tl.log2(sf))) + + # 6) Quantize to FP8 e4m3fn. + y_fp8 = (y / sf[:, None]).to(tl.float8e4nv) + + # 7) Store y and sf. + y_ptrs = y_ptr + offs_m[:, None] * stride_ym + offs_k[None, :] * stride_yn + tl.store(y_ptrs, y_fp8, mask=mask_m[:, None]) + + sf_ptrs = y_sf_ptr + offs_m * stride_sfm + pid_k * stride_sfk + tl.store(sf_ptrs, sf, mask=mask_m) + + +def swiglu_apply_weight_to_fp8_triton( + x: torch.Tensor, + topk_weights: torch.Tensor | None, + clamp_value: float | None = None, + num_per_channels: int = BASELINE_L2_ACT_SF_GRAN, + use_ue8m0_scale: bool = True, +) -> Tuple[torch.Tensor, torch.Tensor]: + """SwiGLU + FP8 quantization. Semantically equivalent to: + gate, up = x[:, :H], x[:, H:] + y = silu(gate.clamp(max=c)) * up.clamp(-c, c) * topk_w + y_sf = y.view(M, H/np, np).abs().amax(-1) / 448 + if use_ue8m0_scale: y_sf = ceil_to_power_of_2(y_sf) + y_fp8 = (y / y_sf.unsqueeze(-1)).to(fp8) + """ + assert x.is_cuda and x.dtype == torch.bfloat16 + assert x.is_contiguous(), "This implementation expects contiguous x" + M, two_H = x.shape + H = two_H // 2 + assert H % num_per_channels == 0, f"H={H} must be divisible by {num_per_channels}" + + y = torch.empty((M, H), dtype=torch.float8_e4m3fn, device=x.device) + y_sf = torch.empty((M, H // num_per_channels), dtype=torch.float32, device=x.device) + + # BLOCK_M=16 keeps register pressure low for the Triton reference kernel. + BLOCK_M = 16 + grid = (triton.cdiv(M, BLOCK_M), H // num_per_channels) + + # Triton still needs a valid pointer when HAS_TOPK=False; x is a placeholder. + topk_ptr = topk_weights if topk_weights is not None else x + + _swiglu_apply_weight_to_fp8_kernel[grid]( + x, + topk_ptr, + y, + y_sf, + M, + H, + x.stride(0), + x.stride(1), + y.stride(0), + y.stride(1), + y_sf.stride(0), + y_sf.stride(1), + float(clamp_value) if clamp_value is not None else 0.0, + HAS_TOPK=topk_weights is not None, + HAS_CLAMP=clamp_value is not None, + USE_UE8M0_SCALE=use_ue8m0_scale, + BLOCK_M=BLOCK_M, + BLOCK_K=num_per_channels, + ) + return y, y_sf + + +# ============================================================================ +# Section 2: grouped weight block-(128, 128) FP8 quantization. +# ---------------------------------------------------------------------------- +# SM90 m_grouped_fp8_gemm_nt_contiguous expects each (128, 128) weight block to +# share one FP32 SF, with K as the inner contiguous SF dimension (K-major). +# Unlike the SM100 FP4 path: +# * deep_gemm.transform_sf_into_required_layout is not needed. +# * SF is FP32, not packed UE8M0. +# ============================================================================ + + +def _quantize_grouped_fp8_block_128_128( + w: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """(G, N, K) bf16 -> (G, N, K) fp8_e4m3fn plus FP32 block SF.""" + g, n, k = w.shape + assert n % 128 == 0 and k % 128 == 0, f"weight N={n}, K={k} must be multiples of 128" + + # Split (N, K) into (N/128, 128, K/128, 128) block interiors. + w_view = w.view(g, n // 128, 128, k // 128, 128).float() + + # In-block absmax -> scale = amax / 448; clamp avoids all-zero scales. + amax = w_view.abs().amax(dim=(-1, -3)).clamp(1e-4) # (G, N/128, K/128) + sf = amax / FP8_E4M3_MAX + + # Divide by the owning block's SF before casting to FP8. + w_fp8 = (w_view / sf.unsqueeze(-1).unsqueeze(-3)).to(torch.float8_e4m3fn) + return w_fp8.view(g, n, k).contiguous(), sf.contiguous() + + +# ============================================================================ +# Section 3: layered accuracy reference and scenarios. +# ============================================================================ + + +def _dequant_block_128_128(w_fp8: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """Inverse of _quantize_grouped_fp8_block_128_128. Returns fp32.""" + *prefix, n, k = w_fp8.shape + assert n % 128 == 0 and k % 128 == 0 + w_view = w_fp8.float().view(*prefix, n // 128, 128, k // 128, 128) + return (w_view * sf.unsqueeze(-1).unsqueeze(-3)).view(*prefix, n, k) + + +def _dequant_per_token_per_128_k(x_fp8: torch.Tensor, sf: torch.Tensor) -> torch.Tensor: + """Dequantize (M, K) fp8 with per-token, per-128-K float scales.""" + m, k = x_fp8.shape + assert k % 128 == 0 + x_view = x_fp8.float().view(m, k // 128, 128) + return (x_view * sf.unsqueeze(-1)).view(m, k) + + +def _swiglu_fp32(gate_up: torch.Tensor, clamp: float) -> torch.Tensor: + """SwiGLU matching the fused SM90 path's clamp semantics.""" + n2 = gate_up.size(-1) + half = n2 // 2 + gate, up = gate_up[..., :half], gate_up[..., half:] + if math.isfinite(clamp): + gate = gate.clamp(max=clamp) + up = up.clamp(min=-clamp, max=clamp) + return torch.nn.functional.silu(gate) * up + + +def _reference_fused( + x_fp8_local: torch.Tensor, + x_sf_local: torch.Tensor, + topk_idx_local: torch.Tensor, + topk_weights_local: torch.Tensor, + l1_w_fp8: torch.Tensor, + l1_w_sf: torch.Tensor, + l2_w_fp8: torch.Tensor, + l2_w_sf: torch.Tensor, + rank_idx: int, + num_ranks: int, + group: dist.ProcessGroup, + num_experts: int, + num_topk: int, + hidden: int, + intermediate_hidden: int, + activation_clamp: float, +) -> torch.Tensor: + """PyTorch BF16/FP32 reference for this rank's fused output.""" + num_experts_per_rank = num_experts // num_ranks + + x_fp8_g = uneven_all_gather(x_fp8_local, group=group) + x_sf_g = uneven_all_gather(x_sf_local, group=group) + topk_idx_g = uneven_all_gather(topk_idx_local, group=group) + topk_w_g = uneven_all_gather(topk_weights_local, group=group) + mg = x_fp8_g.size(0) + + local_size = torch.tensor([x_fp8_local.size(0)], device="cuda", dtype=torch.long) + sizes_t = torch.empty(num_ranks, dtype=torch.long, device="cuda") + dist.all_gather_into_tensor(sizes_t, local_size, group=group) + sizes_list = sizes_t.tolist() + assert sum(sizes_list) == mg + + l1_w_g = [torch.empty_like(l1_w_fp8) for _ in range(num_ranks)] + l1_sf_g = [torch.empty_like(l1_w_sf) for _ in range(num_ranks)] + l2_w_g = [torch.empty_like(l2_w_fp8) for _ in range(num_ranks)] + l2_sf_g = [torch.empty_like(l2_w_sf) for _ in range(num_ranks)] + dist.all_gather(l1_w_g, l1_w_fp8, group=group) + dist.all_gather(l1_sf_g, l1_w_sf, group=group) + dist.all_gather(l2_w_g, l2_w_fp8, group=group) + dist.all_gather(l2_sf_g, l2_w_sf, group=group) + l1_w_all = torch.stack(l1_w_g, dim=0) + l1_sf_all = torch.stack(l1_sf_g, dim=0) + l2_w_all = torch.stack(l2_w_g, dim=0) + l2_sf_all = torch.stack(l2_sf_g, dim=0) + + combine_buf = torch.zeros(mg, num_topk, hidden, dtype=torch.float32, device="cuda") + x_fp32 = _dequant_per_token_per_128_k(x_fp8_g, x_sf_g) + + chunk = 256 + for k in range(num_topk): + mask = topk_idx_g[:, k] >= 0 + if not mask.any(): + continue + sel_idx_full = mask.nonzero(as_tuple=False).squeeze(-1) + for c0 in range(0, sel_idx_full.numel(), chunk): + sel_idx = sel_idx_full[c0:c0 + chunk] + eids = topk_idx_g[sel_idx, k] + weights = topk_w_g[sel_idx, k] + x_sel = x_fp32[sel_idx] + + dst_rank = (eids // num_experts_per_rank).long() + dst_local = (eids % num_experts_per_rank).long() + + l1_w_sel = _dequant_block_128_128( + l1_w_all[dst_rank, dst_local], + l1_sf_all[dst_rank, dst_local], + ) + l1_y = torch.einsum("sk,snk->sn", x_sel, l1_w_sel) + del l1_w_sel + + l1_y = _swiglu_fp32(l1_y, activation_clamp) * weights.unsqueeze(-1) + s, ih = l1_y.shape + assert ih == intermediate_hidden and ih % 64 == 0 + l1_view = l1_y.view(s, ih // 64, 64) + amax = l1_view.abs().amax(dim=-1).clamp(1e-4) + sf2 = amax / FP8_E4M3_MAX + l1_q = (l1_view / sf2.unsqueeze(-1)).to(torch.float8_e4m3fn).float() + l2_in = (l1_q * sf2.unsqueeze(-1)).view(s, ih) + + l2_w_sel = _dequant_block_128_128( + l2_w_all[dst_rank, dst_local], + l2_sf_all[dst_rank, dst_local], + ) + l2_y = torch.einsum("sn,smn->sm", l2_in, l2_w_sel) + del l2_w_sel + + combine_buf[sel_idx, k] = l2_y.to(torch.bfloat16).float() + + y_full_bf16 = combine_buf.to(torch.bfloat16).sum(dim=1).to(torch.bfloat16) + start = sum(sizes_list[:rank_idx]) + end = start + sizes_list[rank_idx] + return y_full_bf16[start:end].contiguous() + + +def _run_accuracy_scenario( + name: str, + cfg: Dict[str, Any], + rank_idx: int, + num_ranks: int, + group: dist.ProcessGroup, + diff_tol: float, +): + num_max = cfg["num_max_tokens_per_rank"] + num_tokens = cfg.get("num_tokens", num_max) + hidden = cfg["hidden"] + intermediate_hidden = cfg["intermediate_hidden"] + num_experts = cfg["num_experts"] + num_topk = cfg["num_topk"] + masked_ratio = cfg.get("masked_ratio", 0.0) + activation_clamp = cfg.get("activation_clamp", 10.0) + fast_math = cfg.get("fast_math", True) + + assert num_experts % num_ranks == 0, ( + f"{name}: experts {num_experts} not divisible by ranks {num_ranks}" + ) + num_experts_per_rank = num_experts // num_ranks + assert num_tokens <= num_max + assert hidden % 128 == 0 and intermediate_hidden % 128 == 0 + + verbose = bool(int(os.environ.get("DG_TEST_VERBOSE", "0"))) + + def trace(stage: str): + if verbose: + print(f"[rank{rank_idx}] {name} :: {stage}", flush=True) + + trace("begin") + torch.manual_seed(rank_idx * 1000 + abs(hash(name)) % 1000) + random.seed(rank_idx * 1000 + abs(hash(name)) % 1000) + + x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + l1_weights_bf16 = torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, + device="cuda", + ) * 0.05 + l2_weights_bf16 = torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + ) * 0.05 + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device="cuda") + topk_weights, topk_idx = torch.topk( + scores, num_topk, dim=-1, largest=True, sorted=False + ) + if masked_ratio > 0: + rand_mask = torch.rand_like(topk_idx, dtype=torch.float) + topk_idx.masked_fill_(rand_mask < masked_ratio, -1) + topk_weights.masked_fill_(topk_idx < 0, 0) + + x_fp8 = per_token_cast_to_fp8( + x_bf16, use_ue8m0=False, gran_k=128, use_packed_ue8m0=False + ) + l1_weights = _quantize_grouped_fp8_block_128_128(l1_weights_bf16) + l2_weights = _quantize_grouped_fp8_block_128_128(l2_weights_bf16) + + trace("weight_transform") + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe_sm90( + l1_weights, l2_weights + ) + + trace("alloc_symm_buffer") + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + num_experts, + num_max, + num_topk, + hidden, + intermediate_hidden, + ) + cum_stats = torch.zeros((num_experts_per_rank,), dtype=torch.int, device="cuda") + + trace("copy_inputs") + buffer.x[:num_tokens].copy_(x_fp8[0]) + buffer.x_sf[:num_tokens].copy_(x_fp8[1]) + buffer.topk_idx[:num_tokens].copy_(topk_idx) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + + y_fused = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + trace("launch_fused") + deep_gemm.fp8_mega_moe( + y_fused, + transformed_l1, + transformed_l2, + buffer, + cumulative_local_expert_recv_stats=cum_stats, + recipe=(128, 128, 128), + activation="swiglu", + activation_clamp=activation_clamp if math.isfinite(activation_clamp) else None, + fast_math=fast_math, + ) + torch.cuda.synchronize() + + trace("reference") + y_ref = _reference_fused( + x_fp8[0], + x_fp8[1], + topk_idx, + topk_weights, + l1_weights[0], + l1_weights[1], + l2_weights[0], + l2_weights[1], + rank_idx, + num_ranks, + group, + num_experts, + num_topk, + hidden, + intermediate_hidden, + activation_clamp, + ) + + diff = calc_diff(y_fused, y_ref) + ok = diff < diff_tol + dist_print( + f" [{name:<32}] diff={diff:.4f} (tol={diff_tol:.2f}) " + f"{'OK' if ok else 'FAIL'}", + once_in_node=True, + ) + assert ok, f"{name}: diff={diff} >= tol={diff_tol}" + if num_tokens > 0 and masked_ratio < 1.0: + assert cum_stats.sum().item() >= 0 + + buffer.destroy() + dist.barrier() + + +_ACCURACY_SMOKE = dict( + num_max_tokens_per_rank=64, + num_tokens=64, + hidden=512, + intermediate_hidden=512, + num_experts=8, + num_topk=2, +) + + +def _accuracy_layer1_smoke() -> List[Tuple[str, Dict[str, Any]]]: + return [("L1.smoke", dict(_ACCURACY_SMOKE))] + + +def _accuracy_layer2_heuristic_branches(num_ranks: int) -> List[Tuple[str, Dict[str, Any]]]: + base = dict( + hidden=1024, + intermediate_hidden=1024, + num_experts=8 * num_ranks, + num_topk=2, + ) + out: List[Tuple[str, Dict[str, Any]]] = [] + for tokens, label in [(64, "small"), (256, "midA"), (512, "midB"), (2048, "large")]: + cfg = dict(base) + cfg.update(num_max_tokens_per_rank=tokens, num_tokens=tokens) + out.append((f"L2.heur.{label}.t{tokens}", cfg)) + return out + + +def _accuracy_layer3_shape_sweep(num_ranks: int) -> List[Tuple[str, Dict[str, Any]]]: + out: List[Tuple[str, Dict[str, Any]]] = [] + base_experts = 8 * num_ranks + for hidden in (512, 2048): + for ih in (512, 2048): + for topk in (1, 2, 4): + if topk > base_experts: + continue + cfg = dict( + num_max_tokens_per_rank=128, + num_tokens=128, + hidden=hidden, + intermediate_hidden=ih, + num_experts=base_experts, + num_topk=topk, + ) + out.append((f"L3.h{hidden}_ih{ih}_k{topk}", cfg)) + return out + + +def _accuracy_layer4_edges(num_ranks: int) -> List[Tuple[str, Dict[str, Any]]]: + base = dict( + num_max_tokens_per_rank=128, + hidden=512, + intermediate_hidden=512, + num_experts=8 * num_ranks, + num_topk=2, + ) + out = [] + for masked_ratio in (0.0, 0.3, 0.7): + cfg = dict(base) + cfg.update(num_tokens=128, masked_ratio=masked_ratio) + out.append((f"L4.mask{masked_ratio:.1f}", cfg)) + cfg = dict(base) + cfg.update(num_tokens=128, masked_ratio=1.0) + out.append(("L4.mask_all", cfg)) + for clamp in (1.0, 10.0, math.inf): + cfg = dict(base) + cfg.update(num_tokens=128, activation_clamp=clamp) + out.append((f"L4.clamp{clamp}", cfg)) + for fast_math in (True, False): + cfg = dict(base) + cfg.update(num_tokens=128, fast_math=fast_math) + out.append((f"L4.fm{int(fast_math)}", cfg)) + cfg = dict(base) + cfg.update(num_tokens=0) + out.append(("L4.tokens0", cfg)) + cfg = dict(base) + cfg.update(num_tokens=base["num_max_tokens_per_rank"]) + out.append(("L4.tokens_max", cfg)) + return out + + +def _accuracy_layer5_stress(num_ranks: int, num_tests: int) -> List[Tuple[str, Dict[str, Any]]]: + rng = random.Random(0xC0FFEE) + out = [] + for i in range(num_tests): + cfg = dict( + num_max_tokens_per_rank=rng.choice([32, 64, 128, 256, 512]), + hidden=rng.choice([512, 1024, 2048]), + intermediate_hidden=rng.choice([512, 1024, 2048]), + num_experts=8 * num_ranks, + num_topk=rng.choice([1, 2, 4]), + masked_ratio=rng.choice([0.0, 0.0, 0.3, 0.5]), + activation_clamp=rng.choice([1.0, 10.0, math.inf]), + fast_math=rng.choice([True, False]), + ) + cfg["num_tokens"] = cfg["num_max_tokens_per_rank"] + out.append((f"L5.rand{i:03d}", cfg)) + return out + + +def _run_accuracy_tests(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + + if get_arch_major() != 9: + dist_print( + f"[SKIP] test_mega_moe_hopper accuracy requires SM90; got SM{get_arch_major()}0", + once_in_node=True, + ) + dist.destroy_process_group() + return + + layers: List[Tuple[str, Dict[str, Any]]] = [] + if 1 in args.layers: + layers += _accuracy_layer1_smoke() + if 2 in args.layers: + layers += _accuracy_layer2_heuristic_branches(num_ranks) + if 3 in args.layers: + layers += _accuracy_layer3_shape_sweep(num_ranks) + if 4 in args.layers: + layers += _accuracy_layer4_edges(num_ranks) + if 5 in args.layers: + layers += _accuracy_layer5_stress(num_ranks, args.num_correctness_tests or 8) + if args.filter: + layers = [(name, cfg) for name, cfg in layers if args.filter in name] + + dist_print( + f"SM90 MegaMoE accuracy plan: {len(layers)} scenarios across " + f"layers {sorted(args.layers)} on {num_ranks} ranks", + once_in_node=True, + ) + + failures: List[str] = [] + for name, cfg in layers: + try: + _run_accuracy_scenario(name, cfg, rank_idx, num_ranks, group, args.diff_tol) + except AssertionError as ex: + dist_print(f" [{name}] FAIL: {ex}", once_in_node=True) + failures.append(name) + if args.fail_fast: + break + + dist_print("", once_in_node=True) + if failures: + dist_print( + f"FAILED {len(failures)}/{len(layers)} scenarios: {failures}", + once_in_node=True, + ) + else: + dist_print(f"PASSED all {len(layers)} scenarios", once_in_node=True) + + dist.barrier() + dist.destroy_process_group() + if failures: + sys.exit(1) + + +# ============================================================================ +# Section 4: optional deep_ep import for dispatch/combine. +# ============================================================================ + + +def _import_deep_ep(): + if _deep_ep is None: + dist_print(f"Failed to import deep_ep: {_deep_ep_import_error}", once_in_node=True) + return None + return _deep_ep + + +class _DeepEPHandle: + def __init__(self, raw_handle, psum_num_recv_tokens_per_expert: torch.Tensor): + self.raw_handle = raw_handle + self.psum_num_recv_tokens_per_expert = psum_num_recv_tokens_per_expert + + +class _DeepEPBufferCompat: + """Compatibility shim for newer DeepEP versions that expose Buffer, not ElasticBuffer.""" + + def __init__(self, deep_ep, group, num_nvl_bytes: int): + self.buffer = deep_ep.Buffer( + group, + num_nvl_bytes=num_nvl_bytes, + num_rdma_bytes=0, + explicitly_destroy=True, + ) + + def dispatch( + self, + x, + *, + topk_idx, + topk_weights, + num_experts: int, + expert_alignment: int, + **_, + ): + num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, event = ( + self.buffer.get_dispatch_layout(topk_idx, num_experts) + ) + recv_x, _, recv_topk_weights, num_recv_tokens_per_expert, raw_handle, event = self.buffer.dispatch( + x, + num_tokens_per_rank=num_tokens_per_rank, + num_tokens_per_rdma_rank=num_tokens_per_rdma_rank, + is_token_in_rank=is_token_in_rank, + num_tokens_per_expert=num_tokens_per_expert, + topk_idx=topk_idx, + topk_weights=topk_weights, + expert_alignment=expert_alignment, + ) + psum = torch.tensor( + num_recv_tokens_per_expert, dtype=torch.int, device=topk_idx.device + ).cumsum(dim=0, dtype=torch.int) + return recv_x, None, recv_topk_weights, _DeepEPHandle(raw_handle, psum), event + + def combine(self, x, *, handle): + raw_handle = handle.raw_handle if isinstance(handle, _DeepEPHandle) else handle + return self.buffer.combine(x, handle=raw_handle) + + def barrier(self, use_comm_stream: bool = False): + torch.cuda.synchronize() + dist.barrier() + + def destroy(self): + self.buffer.destroy() + + +def _make_deep_ep_buffer(deep_ep, group, num_max_tokens_per_rank, hidden, num_topk, sym_buffer_bytes): + if hasattr(deep_ep, "ElasticBuffer"): + return deep_ep.ElasticBuffer( + group, + num_max_tokens_per_rank=num_max_tokens_per_rank, + hidden=hidden, + num_topk=num_topk, + use_fp8_dispatch=True, + explicitly_destroy=True, + allow_multiple_reduction=False, + ) + nvl_alignment = 2 * 1024 * 1024 + num_nvl_bytes = ((int(sym_buffer_bytes) + nvl_alignment - 1) // nvl_alignment) * nvl_alignment + return _DeepEPBufferCompat(deep_ep, group, num_nvl_bytes=num_nvl_bytes) + + +def _make_deep_ep_low_latency_buffer( + deep_ep, group, num_max_dispatch_tokens_per_rank, hidden, num_experts +): + """Build a DeepEP ``Buffer`` configured for low-latency dispatch/combine. + + Mirrors the buffer construction used by sglang's + ``_DeepEPDispatcherImplLowLatency`` (see + ``sglang/srt/layers/moe/token_dispatcher/deepep.py``): RDMA bytes from + ``get_low_latency_rdma_size_hint`` and one QP per local expert. + """ + num_rdma_bytes = deep_ep.Buffer.get_low_latency_rdma_size_hint( + num_max_dispatch_tokens_per_rank, hidden, group.size(), num_experts + ) + return deep_ep.Buffer( + group, + num_nvl_bytes=0, + num_rdma_bytes=num_rdma_bytes, + low_latency_mode=True, + num_qps_per_rank=num_experts // group.size(), + allow_nvlink_for_low_latency_mode=True, + explicitly_destroy=True, + ) + + +# ---------------------------------------------------------------------------- +# Masked SwiGLU + FP8 quantization (low-latency layout). +# ---------------------------------------------------------------------------- +# DeepEP low-latency dispatch returns tokens packed as +# x: [num_local_experts, num_max_dispatch_tokens_per_rank * num_ranks, 2*IH] +# with a per-expert valid-token count ``masked_m[g]``. The masked GEMM does +# not care about the trailing junk rows, but to feed the L2 masked GEMM with +# correct scales we still need a per-token-per-128-K FP32 scale tensor with +# the same masked-layout convention. This kernel produces: +# y : [E, M, IH] fp8_e4m3fn +# y_sf : [E, M, IH // BLOCK_K] fp32 (row-major; DeepGEMM SM90 path +# accepts this layout) +# Modeled on sglang's ``_silu_and_mul_post_quant_kernel`` in +# ``sglang/srt/layers/moe/ep_moe/kernels.py``. + + +@triton.jit +def _swiglu_masked_post_quant_kernel( + x_ptr, + stride_x_e, + stride_x_m, + stride_x_n, + y_ptr, + stride_y_e, + stride_y_m, + stride_y_n, + y_sf_ptr, + stride_sf_e, + stride_sf_m, + stride_sf_k, + masked_m_ptr, + H, + clamp_value, + HAS_CLAMP: tl.constexpr, + USE_UE8M0_SCALE: tl.constexpr, + BLOCK_K: tl.constexpr, + NUM_STAGES: tl.constexpr, +): + pid_k = tl.program_id(0) # column tile within IH + pid_m = tl.program_id(1) # token-stripe within this expert + pid_e = tl.program_id(2) # expert + + num_token_stripes = tl.num_programs(1) + num_valid_tokens = tl.load(masked_m_ptr + pid_e) + + offs_k = pid_k * BLOCK_K + tl.arange(0, BLOCK_K) + + # Element ptrs for one (expert, token_index, k_block). + x_base = x_ptr + pid_e * stride_x_e + offs_k * stride_x_n + y_base = y_ptr + pid_e * stride_y_e + offs_k * stride_y_n + sf_base = y_sf_ptr + pid_e * stride_sf_e + pid_k * stride_sf_k + + for token in tl.range(pid_m, num_valid_tokens, num_token_stripes, num_stages=NUM_STAGES): + gate = tl.load(x_base + token * stride_x_m).to(tl.float32) + up = tl.load(x_base + token * stride_x_m + H * stride_x_n).to(tl.float32) + + if HAS_CLAMP: + gate = tl.minimum(gate, clamp_value) + up = tl.minimum(tl.maximum(up, -clamp_value), clamp_value) + + y = gate * tl.sigmoid(gate) * up + + amax = tl.max(tl.abs(y)) + sf = tl.maximum(amax / _FP8_E4M3_MAX_TL, 1.0e-30) + if USE_UE8M0_SCALE: + sf = tl.exp2(tl.ceil(tl.log2(sf))) + + y_fp8 = (y / sf).to(tl.float8e4nv) + + tl.store(y_base + token * stride_y_m, y_fp8) + tl.store(sf_base + token * stride_sf_m, sf) + + +def swiglu_masked_post_quant_to_fp8( + x: torch.Tensor, + masked_m: torch.Tensor, + quant_group_size: int = BASELINE_L2_ACT_SF_GRAN, + clamp_value: float | None = None, + use_ue8m0_scale: bool = False, +) -> Tuple[torch.Tensor, torch.Tensor]: + """SwiGLU + per-(token, BLOCK_K) FP8 quant on masked-layout input. + + Input: + x : (E, M, 2*H) bf16, contiguous + masked_m : (E,) int, number of valid rows per expert + Returns: + y : (E, M, H) fp8_e4m3fn + y_sf : (E, M, H // quant_group_size) fp32 (row-major) + + The MoE low-latency path applies topk weights inside + ``low_latency_combine``, so this kernel does NOT multiply by topk weights. + """ + assert x.is_cuda and x.dtype == torch.bfloat16 + assert x.is_contiguous(), "Expects contiguous masked-layout input" + assert x.dim() == 3 and x.shape[-1] % 2 == 0 + E, M, two_H = x.shape + H = two_H // 2 + assert H % quant_group_size == 0 + assert masked_m.shape == (E,) + + y = torch.empty((E, M, H), dtype=torch.float8_e4m3fn, device=x.device) + y_sf = torch.empty( + (E, M, H // quant_group_size), dtype=torch.float32, device=x.device + ) + + BLOCK_K = quant_group_size + # Heuristic similar to sglang's silu_and_mul_masked_post_quant_fwd. + block_num_per_expert = 64 if E < 4 else 32 + + grid = (H // BLOCK_K, block_num_per_expert, E) + + _swiglu_masked_post_quant_kernel[grid]( + x, + x.stride(0), + x.stride(1), + x.stride(2), + y, + y.stride(0), + y.stride(1), + y.stride(2), + y_sf, + y_sf.stride(0), + y_sf.stride(1), + y_sf.stride(2), + masked_m, + H, + float(clamp_value) if clamp_value is not None else 0.0, + HAS_CLAMP=clamp_value is not None, + USE_UE8M0_SCALE=use_ue8m0_scale, + BLOCK_K=BLOCK_K, + NUM_STAGES=4, + num_warps=1, + ) + return y, y_sf + + +# ============================================================================ +# Section 5: CUDA event median timing, independent of tilelang.do_bench. +# ============================================================================ + + +def _bench_cuda_events( + fn, num_warmup: int = 5, num_repeat: int = 20, l2_flush_gb: float = 8.0 +) -> float: + """Return median runtime of fn in seconds.""" + for _ in range(num_warmup): + fn() + torch.cuda.synchronize() + times_ms = [] + for _ in range(num_repeat): + # Flush L2 to avoid optimistic timings from repeated cache hits. + if l2_flush_gb > 0: + free_bytes, _ = torch.cuda.mem_get_info() + flush_bytes = min(int(l2_flush_gb * 1e9), int(free_bytes * 0.5)) + if flush_bytes >= 4: + torch.empty(flush_bytes // 4, dtype=torch.int, device="cuda").zero_() + s = torch.cuda.Event(enable_timing=True) + e = torch.cuda.Event(enable_timing=True) + s.record() + fn() + e.record() + e.synchronize() + times_ms.append(s.elapsed_time(e)) + times_ms.sort() + return times_ms[len(times_ms) // 2] / 1e3 + + +# ============================================================================ +# Section 6: fused-only sweep / NCU profile mode. +# ============================================================================ + + +def _run_fused_only_config( + args: argparse.Namespace, + num_tokens: int, + num_max_tokens_per_rank: int, + hidden: int, + intermediate_hidden: int, + num_experts: int, + num_topk: int, + num_ranks: int, + rank_idx: int, + group: dist.ProcessGroup, +): + num_experts_per_rank = num_experts // num_ranks + assert num_tokens <= num_max_tokens_per_rank + assert num_experts % num_ranks == 0, ( + f"num_experts={num_experts} must be divisible by num_ranks={num_ranks}" + ) + assert hidden % 128 == 0 + assert intermediate_hidden % 128 == 0 + assert intermediate_hidden // 64 <= 64, ( + f"SM90 fused kernel requires intermediate_hidden <= 4096, got {intermediate_hidden}" + ) + + buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + ) + + x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + l1_weights_bf16 = torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, + device="cuda", + ) * 0.05 + l2_weights_bf16 = torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + ) * 0.05 + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device="cuda") + topk_weights, topk_idx = torch.topk( + scores, num_topk, dim=-1, largest=True, sorted=False + ) + if args.masked_ratio > 0: + rand_mask = torch.rand_like(topk_idx, dtype=torch.float) + topk_idx.masked_fill_(rand_mask < args.masked_ratio, -1) + topk_weights.masked_fill_(topk_idx < 0, 0) + + x_fp8 = per_token_cast_to_fp8( + x_bf16, use_ue8m0=False, gran_k=128, use_packed_ue8m0=False + ) + l1_weights = _quantize_grouped_fp8_block_128_128(l1_weights_bf16) + l2_weights = _quantize_grouped_fp8_block_128_128(l2_weights_bf16) + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe_sm90( + l1_weights, l2_weights + ) + + cum_stats = torch.zeros((num_experts_per_rank,), dtype=torch.int, device="cuda") + y_fused = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + clamp_arg = args.activation_clamp if math.isfinite(args.activation_clamp) else None + + def run_fused(): + buffer.x[:num_tokens].copy_(x_fp8[0]) + buffer.x_sf[:num_tokens].copy_(x_fp8[1]) + buffer.topk_idx[:num_tokens].copy_(topk_idx) + buffer.topk_weights[:num_tokens].copy_(topk_weights) + deep_gemm.fp8_mega_moe( + y_fused, + transformed_l1, + transformed_l2, + buffer, + cumulative_local_expert_recv_stats=cum_stats, + recipe=(128, 128, 128), + activation="swiglu", + activation_clamp=clamp_arg, + fast_math=bool(args.fast_math), + ) + return y_fused + + if args.ncu_profile_only: + dist_print( + f"[NCU] tokens={num_tokens} hidden={hidden} ih={intermediate_hidden}", + once_in_node=True, + ) + run_fused() + torch.cuda.synchronize() + dist.barrier() + buffer.destroy() + return + + run_fused() + dist.barrier() + t_fused = bench_kineto( + run_fused, + SM90_KERNEL_NAME, + barrier=lambda: dist.barrier(), + num_tests=args.num_bench_tests, + suppress_kineto_output=True, + ) + + gathered_topk_idx = uneven_all_gather(topk_idx, group=group) + gathered_topk_idx[ + (gathered_topk_idx < rank_idx * num_experts_per_rank) + | (gathered_topk_idx >= (rank_idx + 1) * num_experts_per_rank) + ] = -1 + local_expert_ids = gathered_topk_idx[gathered_topk_idx != -1] + num_recv_tokens = int(local_expert_ids.numel()) + num_touched_experts = int(torch.unique(local_expert_ids).numel()) if num_recv_tokens else 0 + + def safe_div(a, b): + return float("nan") if b == 0 else a / b + + tflops = safe_div( + 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, + t_fused, + ) + num_hbm_bytes = ( + num_touched_experts * intermediate_hidden * 2 * hidden + + num_touched_experts * hidden * intermediate_hidden + + num_recv_tokens * hidden + + num_recv_tokens * intermediate_hidden + + num_recv_tokens * intermediate_hidden + + num_recv_tokens * hidden * 2 + ) + hbm_gbs = safe_div(num_hbm_bytes / 1e9, t_fused) + + dist_print( + f" tokens={num_tokens:4d} recv={num_recv_tokens:5d} " + f"experts={num_touched_experts:4d} {t_fused * 1e6:7.1f} us " + f"{tflops:6.1f} TFLOPS {hbm_gbs:6.0f} GB/s (rank{rank_idx})", + once_in_node=True, + ) + + dist.barrier() + buffer.destroy() + + +def _run_fused_only_sweep(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + torch.manual_seed(rank_idx) + random.seed(rank_idx) + + if get_arch_major() != 9: + dist_print( + f"[SKIP] test_mega_moe_hopper fused-only sweep requires SM90; got SM{get_arch_major()}0", + once_in_node=True, + ) + dist.destroy_process_group() + return + + batches = args.batches if args.batches is not None else [1, 2, 4, 8, 16, 32] + if args.ncu_profile_only: + batches = batches[:1] + + dist_print( + f"SM90 MegaMoE fused-only sweep: ranks={num_ranks} hidden={args.hidden} " + f"ih={args.intermediate_hidden} experts={args.num_experts} topk={args.num_topk} " + f"masked_ratio={args.masked_ratio} fast_math={bool(args.fast_math)}", + once_in_node=True, + ) + + num_max_tokens_per_rank = max(batches) + for num_tokens in batches: + _run_fused_only_config( + args, + num_tokens, + num_max_tokens_per_rank, + args.hidden, + args.intermediate_hidden, + args.num_experts, + args.num_topk, + num_ranks, + rank_idx, + group, + ) + + dist.barrier() + dist.destroy_process_group() + + +# ============================================================================ +# Section 7: per-rank comparison entry point. +# ============================================================================ + + +def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): + if args.accuracy: + _run_accuracy_tests(local_rank, num_local_ranks, args) + return + + if args.fused_only_sweep: + _run_fused_only_sweep(local_rank, num_local_ranks, args) + return + + # Initialize distributed state; rank_idx is global rank and group is NCCL. + rank_idx, num_ranks, group = init_dist(local_rank, num_local_ranks) + torch.manual_seed(rank_idx) + random.seed(rank_idx) + + if get_arch_major() != 9: + dist_print( + f"[SKIP] test_mega_moe_hopper requires SM90; got SM{get_arch_major()}0", + once_in_node=True, + ) + dist.destroy_process_group() + return + + # Shape parameters, with names matching tests/test_mega_moe.py. + num_max_tokens_per_rank = args.num_max_tokens_per_rank + num_tokens = ( + max( + 0, + args.num_max_tokens_per_rank + - random.randint(0, args.num_max_removed_tokens), + ) + if args.num_tokens == 0 + else args.num_tokens + ) + hidden, intermediate_hidden = args.hidden, args.intermediate_hidden + num_experts, num_topk = args.num_experts, args.num_topk + num_experts_per_rank = num_experts // num_ranks + assert num_tokens <= num_max_tokens_per_rank + assert num_experts % num_ranks == 0, ( + f"num_experts={num_experts} must be divisible by num_ranks={num_ranks}" + ) + + # SM90 fused-kernel shape constraints from csrc/apis/sm90_mega.hpp::fp8_mega_moe: + # * H and IH must be multiples of 128 (L1 input per-128-K SF and + # block-(128,128) weight SF). + # * IH / 64 <= 64, i.e. IH <= 4096, because l2_arrival_mask is uint64 + # with one bit per 64-column block. + assert hidden % 128 == 0 + assert intermediate_hidden % 128 == 0 + assert intermediate_hidden // 64 <= 64, ( + f"SM90 fused kernel requires intermediate_hidden <= 4096, got {intermediate_hidden}" + ) + + # ---- Create BF16 token and weight inputs ---- + # x: local tokens for this rank. + x_bf16 = torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + # L1 weight maps hidden -> 2*intermediate_hidden (gate and up packed). + l1_weights_bf16 = torch.randn( + (num_experts_per_rank, intermediate_hidden * 2, hidden), + dtype=torch.bfloat16, + device="cuda", + ) + # L2 weight maps intermediate_hidden -> hidden. + l2_weights_bf16 = torch.randn( + (num_experts_per_rank, hidden, intermediate_hidden), + dtype=torch.bfloat16, + device="cuda", + ) + + # Routing: scores -> topk_idx (M, K) and topk_weights (M, K). + scores = torch.randn((num_tokens, num_experts), dtype=torch.float, device="cuda") + topk_weights, topk_idx = torch.topk( + scores, num_topk, dim=-1, largest=True, sorted=False + ) + if args.masked_ratio > 0: + rand_mask = torch.rand_like(topk_idx, dtype=torch.float) + topk_idx.masked_fill_(rand_mask < args.masked_ratio, -1) + topk_weights.masked_fill_(topk_idx < 0, 0) + + # Keep separate recv counters so fused and baseline do not overwrite each other. + cum_stats_fused = torch.zeros( + (num_experts_per_rank,), dtype=torch.int, device="cuda" + ) + cum_stats_baseline = cum_stats_fused.clone() + + # ---- BF16 -> FP8 quantization ---- + # x_fp8 is (token_fp8 (M, hidden), token_sf (M, hidden//128) row-major FP32). + # SM90 expects FP32 SF, not packed UE8M0. + x_fp8 = per_token_cast_to_fp8( + x_bf16, use_ue8m0=False, gran_k=128, use_packed_ue8m0=False + ) + + # Weight quantization: (G, N, K) bf16 -> FP8 e4m3fn plus block FP32 SF. + # The DeepEP grouped-GEMM baseline uses these untransformed tuples directly. + l1_weights = _quantize_grouped_fp8_block_128_128(l1_weights_bf16) + l2_weights = _quantize_grouped_fp8_block_128_128(l2_weights_bf16) + + # Fused path: interleave gate/up along N for FP8 L1 weights; SF is unchanged. + transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe_sm90( + l1_weights, l2_weights + ) + + # SwiGLU clamp: finite values enable clamp; inf maps to None and disables it. + clamp_arg = args.activation_clamp if math.isfinite(args.activation_clamp) else None + run_baseline_enabled = args.run_baseline or bool(args.check_output_diff) + run_ll_baseline_enabled = bool(args.run_low_latency_baseline) + + # ---- M-dimension alignment for the grouped-GEMM baseline ---- + alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout() + deep_gemm.set_mk_alignment_for_contiguous_layout(alignment) + + # ---- Allocate fused SymmBuffer and output buffer ---- + sym_buffer = deep_gemm.get_symm_buffer_for_mega_moe( + group, + num_experts, + num_max_tokens_per_rank, + num_topk, + hidden, + intermediate_hidden, + ) + y_fused = torch.empty((num_tokens, hidden), dtype=torch.bfloat16, device="cuda") + + def run_fused(): + # Match the SM100 test: DG_COMM_KERNEL_DEBUG=1 zeros the whole + # sym_buffer at kernel exit, so inputs must be re-copied every call. + sym_buffer.x[:num_tokens].copy_(x_fp8[0]) + sym_buffer.x_sf[:num_tokens].copy_(x_fp8[1]) + sym_buffer.topk_idx[:num_tokens].copy_(topk_idx) + sym_buffer.topk_weights[:num_tokens].copy_(topk_weights) + + deep_gemm.fp8_mega_moe( + y_fused, + transformed_l1, + transformed_l2, + sym_buffer, + cumulative_local_expert_recv_stats=cum_stats_fused, + recipe=(128, 128, 128), + activation="swiglu", + activation_clamp=clamp_arg, + fast_math=bool(args.fast_math), + ) + return y_fused + + # ---- Print config ---- + dist_print("Config (SM90 fused MegaMoE):", once_in_node=True) + dist_print(f" > Tokens: {num_tokens}/{num_max_tokens_per_rank}", once_in_node=True) + dist_print( + f" > Hidden: {hidden}, Intermediate: {intermediate_hidden}", once_in_node=True + ) + dist_print( + f" > Experts: {num_topk}/{num_experts} (per-rank: {num_experts_per_rank})", + once_in_node=True, + ) + dist_print(f" > Masked ratio: {args.masked_ratio}", once_in_node=True) + dist_print( + f" > Activation SF: fused L2 per-{FUSED_L2_ACT_SF_GRAN} FP32 pow2, " + f"baseline L2 per-{BASELINE_L2_ACT_SF_GRAN} FP32 pow2 " + f"(SM90 grouped-GEMM constraint)", + once_in_node=True, + ) + dist_print( + f" > Baseline: {'enabled' if run_baseline_enabled else 'disabled'}", + once_in_node=True, + ) + dist_print( + f" > Low-latency baseline: {'enabled' if run_ll_baseline_enabled else 'disabled'}", + once_in_node=True, + ) + dist_print( + f" > Buffer: {sym_buffer.buffer.nbytes / 2**30:.3f} GiB", once_in_node=True + ) + dist_print(once_in_node=True) + + # Match tests/test_mega_moe.py: NCU mode runs only the fused kernel to avoid + # baseline noise in the profile. + if args.ncu_profile_only: + dist_print("Run fused SM90 mega-MoE kernel:", once_in_node=True) + y = run_fused() + torch.cuda.synchronize() + assert y.shape == (num_tokens, hidden) and y.dtype == torch.bfloat16 + dist_print(" > Done, exiting", once_in_node=True) + dist.barrier() + sym_buffer.destroy() + dist.destroy_process_group() + return + + # ---- Allocate DeepEP buffer for the baseline ---- + deep_ep = ( + _import_deep_ep() + if (run_baseline_enabled or run_ll_baseline_enabled) + else None + ) + ep_buffer = None + if deep_ep is not None and run_baseline_enabled: + ep_buffer = _make_deep_ep_buffer( + deep_ep, + group, + num_max_tokens_per_rank, + hidden, + num_topk, + sym_buffer.buffer.nbytes, + ) + + # ---- Allocate DeepEP buffer for the low-latency baseline ---- + # Low-latency mode requires its own ``Buffer(low_latency_mode=True, ...)`` + # with ``num_qps_per_rank == num_local_experts`` and RDMA bytes sized via + # ``get_low_latency_rdma_size_hint``. See sglang's + # ``DeepEPBuffer.get_deepep_buffer`` for the canonical setup. + ll_buffer = None + if deep_ep is not None and run_ll_baseline_enabled: + ll_buffer = _make_deep_ep_low_latency_buffer( + deep_ep, + group, + num_max_tokens_per_rank, + hidden, + num_experts, + ) + + # ---------------------------------------------------------------- + # Baseline body: dispatch -> L1 GEMM -> SwiGLU+quantize -> L2 GEMM -> combine. + # It uses the same FP8 weights and FP32 block-(128,128) SF as the fused path, + # but without the fused-only gate/up interleave. + # ---------------------------------------------------------------- + def run_baseline(): + recv_x, _, recv_topk_weights, handle, _ = ep_buffer.dispatch( + x_fp8, + topk_idx=topk_idx, + topk_weights=topk_weights, + cumulative_local_expert_recv_stats=cum_stats_baseline, + num_experts=num_experts, + expert_alignment=alignment, + do_cpu_sync=False, + do_handle_copy=False, + do_expand=True, + use_tma_aligned_col_major_sf=False, # SM90: row-major float SF + ) + n = recv_x[0].size(0) + + # L1 GEMM: FP8 token @ FP8 W1 -> BF16 intermediate activation (gate||up). + l1_y = torch.empty( + (n, intermediate_hidden * 2), dtype=torch.bfloat16, device="cuda" + ) + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + recv_x, + l1_weights, + l1_y, + handle.psum_num_recv_tokens_per_expert, + use_psum_layout=True, + disable_ue8m0_cast=True, + ) + + # Triton SwiGLU + FP8 quantization, including topk weight scaling. + # The fused SM90 MegaMoE L2 activation SF is per-64-K. The current + # DeepGEMM SM90 grouped GEMM supports only per-128-K activation SF, so + # the baseline uses per-128-K FP32 scales with the same power-of-two + # rounding rule as the fused epilogue. + l1_y = swiglu_apply_weight_to_fp8_triton( + x=l1_y, + topk_weights=recv_topk_weights, + clamp_value=clamp_arg, + num_per_channels=BASELINE_L2_ACT_SF_GRAN, + use_ue8m0_scale=True, + ) + + # L2 GEMM: FP8 intermediate activation @ FP8 W2 -> BF16. + l2_y = torch.empty((n, hidden), dtype=torch.bfloat16, device="cuda") + deep_gemm.m_grouped_fp8_gemm_nt_contiguous( + l1_y, + l2_weights, + l2_y, + handle.psum_num_recv_tokens_per_expert, + use_psum_layout=True, + disable_ue8m0_cast=True, + ) + + # DeepEP combine: gather each token's topk expert outputs back to source rank. + return ep_buffer.combine(l2_y, handle=handle)[0] + + # ---------------------------------------------------------------- + # Low-latency baseline body. Mirrors the sglang + # ``_DeepEPDispatcherImplLowLatency`` pipeline: + # 1. ``low_latency_dispatch(use_fp8=True)`` -> per-expert packed FP8 tokens + # with shape ``[E_local, M_max, hidden]`` plus FP32 scales + # ``[E_local, M_max, hidden // 128]`` and per-expert ``masked_m``. + # 2. Masked grouped FP8 GEMM with L1 weights. + # 3. Masked SwiGLU + per-128-K FP8 quantize. (topk weights are NOT + # applied here — ``low_latency_combine`` applies them internally.) + # 4. Masked grouped FP8 GEMM with L2 weights. + # 5. ``low_latency_combine`` (reduces with topk weights). + # ---------------------------------------------------------------- + if run_ll_baseline_enabled: + M_max_ll = num_max_tokens_per_rank * num_ranks + # Expected per-expert mean of ``masked_m`` after dispatch. Used as the + # ``expected_m`` hint for the DeepGEMM masked kernel selector. + expected_m_ll = max( + 1, + (num_max_tokens_per_rank * num_ranks * num_topk + num_experts - 1) + // num_experts, + ) + # Pre-allocate per-call output buffers; the masked GEMM writes into them + # in place and ignores rows past ``masked_m[g]``. + ll_l1_y = torch.empty( + (num_experts_per_rank, M_max_ll, intermediate_hidden * 2), + dtype=torch.bfloat16, + device="cuda", + ) + ll_l2_y = torch.empty( + (num_experts_per_rank, M_max_ll, hidden), + dtype=torch.bfloat16, + device="cuda", + ) + ll_combined = torch.empty( + (num_tokens, hidden), dtype=torch.bfloat16, device="cuda" + ) + # DeepEP low-latency dispatch requires int64 topk indices. + topk_idx_ll = topk_idx.to(torch.int64) + + def run_baseline_low_latency(): + # 1) Low-latency dispatch with FP8 cast. + (recv_x_data, recv_x_sf), masked_m, ll_handle, event, hook = ( + ll_buffer.low_latency_dispatch( + x_bf16, + topk_idx_ll, + num_max_tokens_per_rank, + num_experts, + use_fp8=True, + round_scale=False, + use_ue8m0=False, + async_finish=False, + return_recv_hook=False, + ) + ) + + # 2) L1 masked grouped FP8 GEMM: + # (E_local, M_max, hidden) @ (E_local, 2*IH, hidden)^T -> (E_local, M_max, 2*IH). + deep_gemm.fp8_m_grouped_gemm_nt_masked( + (recv_x_data, recv_x_sf), + l1_weights, + ll_l1_y, + masked_m, + expected_m_ll, + disable_ue8m0_cast=True, + ) + + # 3) Masked SwiGLU + per-128-K FP8 quant. Topk weights are NOT applied + # here — they are reduced inside ``low_latency_combine``. + l1_fp8, l1_sf = swiglu_masked_post_quant_to_fp8( + ll_l1_y, + masked_m, + quant_group_size=BASELINE_L2_ACT_SF_GRAN, + clamp_value=clamp_arg, + use_ue8m0_scale=False, + ) + + # 4) L2 masked grouped FP8 GEMM: + # (E_local, M_max, IH) @ (E_local, H, IH)^T -> (E_local, M_max, H). + deep_gemm.fp8_m_grouped_gemm_nt_masked( + (l1_fp8, l1_sf), + l2_weights, + ll_l2_y, + masked_m, + expected_m_ll, + disable_ue8m0_cast=True, + ) + + # 5) Low-latency combine: per-token weighted reduction across topk + # expert replicas; outputs (num_tokens, hidden) bf16. + combined_x, event, hook = ll_buffer.low_latency_combine( + ll_l2_y, + topk_idx_ll, + topk_weights, + ll_handle, + use_logfmt=False, + zero_copy=False, + async_finish=False, + return_recv_hook=False, + out=ll_combined, + ) + return combined_x + + # ---- Run once to check fused and optional baseline paths ---- + y = run_fused() + assert y.shape == (num_tokens, hidden) and y.dtype == torch.bfloat16, ( + f"unexpected fused output shape/dtype: shape={y.shape}, dtype={y.dtype}" + ) + if ep_buffer is not None: + out_b = run_baseline() + assert out_b.shape == (num_tokens, hidden) and out_b.dtype == torch.bfloat16, ( + f"unexpected baseline output shape/dtype: shape={out_b.shape}, dtype={out_b.dtype}" + ) + if args.check_output_diff: + diff = (y.float() - out_b.float()).abs() + denom = out_b.float().abs().mean().clamp_min(1e-12) + dist_print( + "Output diff (fused vs per-128 baseline):", once_in_node=True + ) + dist_print( + f" > max_abs={diff.max().item():.6e}, " + f"mean_abs={diff.mean().item():.6e}, " + f"mean_abs/mean_ref={diff.mean().div(denom).item():.6e}", + once_in_node=True, + ) + dist_print(once_in_node=True) + if ll_buffer is not None: + out_ll = run_baseline_low_latency() + assert out_ll.shape == (num_tokens, hidden) and out_ll.dtype == torch.bfloat16, ( + f"unexpected LL baseline output shape/dtype: shape={out_ll.shape}, dtype={out_ll.dtype}" + ) + if args.check_output_diff: + diff = (y.float() - out_ll.float()).abs() + denom = out_ll.float().abs().mean().clamp_min(1e-12) + dist_print( + "Output diff (fused vs low-latency baseline):", once_in_node=True + ) + dist_print( + f" > max_abs={diff.max().item():.6e}, " + f"mean_abs={diff.mean().item():.6e}, " + f"mean_abs/mean_ref={diff.mean().div(denom).item():.6e}", + once_in_node=True, + ) + dist_print(once_in_node=True) + + # ---- Count tokens routed to this rank and touched local experts ---- + # Gather all topk_idx tensors and mark entries outside this rank's local + # expert range as -1. Remaining entries are routed (token, slot) pairs. + gathered_topk_idx = uneven_all_gather(topk_idx, group=group) + gathered_topk_idx[ + (gathered_topk_idx < rank_idx * num_experts_per_rank) + | (gathered_topk_idx >= (rank_idx + 1) * num_experts_per_rank) + ] = -1 + local_expert_ids = gathered_topk_idx[gathered_topk_idx != -1] + num_recv_tokens = int(local_expert_ids.numel()) + num_touched_experts = int(torch.unique(local_expert_ids).numel()) + + # ---- benchmark ---- + # Fused: bench_kineto selects the sm90_fp8_mega_moe_impl GPU region only. + t_fused = bench_kineto( + run_fused, + SM90_KERNEL_NAME, + num_tests=args.num_bench_tests, + barrier=lambda: ep_buffer.barrier(use_comm_stream=False) + if ep_buffer is not None + else dist.barrier(), + trace_path=( + f"{args.dump_profile_traces}/mega_moe_hopper_rank{rank_idx}.json" + if args.dump_profile_traces + else None + ), + ) + # Baseline: use CUDA event median timing for consistency across SM90 setups. + t_baseline = ( + _bench_cuda_events( + run_baseline, + num_warmup=args.num_warmup, + num_repeat=args.num_repeat, + l2_flush_gb=args.l2_flush_gb, + ) + if ep_buffer is not None + else 0.0 + ) + # Low-latency baseline timing (same CUDA-event median methodology). + t_baseline_ll = ( + _bench_cuda_events( + run_baseline_low_latency, + num_warmup=args.num_warmup, + num_repeat=args.num_repeat, + l2_flush_gb=args.l2_flush_gb, + ) + if ll_buffer is not None + else 0.0 + ) + + def safe_div(a, b): + return float("nan") if b == 0 else a / b + + # End-to-end TFLOPS: three matmuls (L1 gate, L1 up, L2), each 2*M*N*K. + tflops = safe_div( + 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_fused + ) + + # HBM byte estimate (SM90 weights are FP8 = 1B/elem, unlike SM100 FP4). + l1_weight_bytes = num_touched_experts * intermediate_hidden * 2 * hidden + l2_weight_bytes = num_touched_experts * hidden * intermediate_hidden + l1_weight_sf_bytes = ( + num_touched_experts + * (intermediate_hidden * 2 // WEIGHT_SF_GRAN_MN) + * (hidden // WEIGHT_SF_GRAN_K) + * 4 + ) + l2_weight_sf_bytes = ( + num_touched_experts + * (hidden // WEIGHT_SF_GRAN_MN) + * (intermediate_hidden // WEIGHT_SF_GRAN_K) + * 4 + ) + l1_input_sf_bytes = num_recv_tokens * (hidden // L1_ACT_SF_GRAN) * 4 + l2_act_sf_bytes = ( + num_recv_tokens * (intermediate_hidden // FUSED_L2_ACT_SF_GRAN) * 4 + ) + num_hbm_bytes = ( + l1_weight_bytes + + l2_weight_bytes # weights (FP8) + + l1_weight_sf_bytes + + l2_weight_sf_bytes # weight SF (FP32) + + num_recv_tokens * hidden + + l1_input_sf_bytes # L1 input read (FP8 + SF) + + num_recv_tokens * intermediate_hidden + + l2_act_sf_bytes # L1 output write (FP8 + SF) + + num_recv_tokens * intermediate_hidden + + l2_act_sf_bytes # L2 input read (FP8 + SF) + + num_recv_tokens * hidden * 2 # L2 output write (BF16) + ) + hbm_gbs = safe_div(num_hbm_bytes / 1e9, t_fused) + + # NVLink bytes: dispatch pulls token + input SF + topk weight; combine writes BF16. + num_nvlink_bytes = num_recv_tokens * (hidden + hidden // 32 + 4 + hidden * 2) + nvlink_gbs = safe_div(num_nvlink_bytes / 1e9, t_fused) + + # Serial lower bound for combine reduction, using 6.5e12 B/s as an estimate. + t_reduction = num_tokens * hidden * 2 * (1 + num_topk) / 6.5e12 + + # Overlap adjustment: remove the non-overlapped serial reduction estimate. + approx_factor = t_fused / max(t_fused - t_reduction, 1e-12) + + # Baseline uses the same FLOPs and HBM byte estimate, with t_baseline. + tflops_baseline = safe_div( + 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_baseline + ) + hbm_gbs_baseline = safe_div(num_hbm_bytes / 1e9, t_baseline) + nvlink_gbs_baseline = safe_div(num_nvlink_bytes / 1e9, t_baseline) + # Low-latency baseline pads each expert's activation to ``M_max_ll``, so + # the weights are streamed once per expert regardless of routing. NVLink + # bytes match the normal-mode baseline (same per-routed-token volume). + tflops_baseline_ll = safe_div( + 2 * num_recv_tokens * (hidden * intermediate_hidden * 3) / 1e12, t_baseline_ll + ) + hbm_gbs_baseline_ll = safe_div(num_hbm_bytes / 1e9, t_baseline_ll) + nvlink_gbs_baseline_ll = safe_div(num_nvlink_bytes / 1e9, t_baseline_ll) + + def fmt_perf_line( + name: str, + t: float, + compute_tflops: float, + hbm_gbs_: float, + nvlink_gbs_: float, + reduction_us: float | None = None, + speedup: float | None = None, + ) -> str: + reduction = f"{reduction_us:13.1f}" if reduction_us is not None else f"{'-':>13}" + speedup_text = ( + f"{speedup:6.2f}x {'fused faster' if speedup > 1 else 'baseline faster'}" + if speedup is not None else + f"{'-':>21}" + ) + return ( + f" > {name:<10} {rank_idx:2d}/{num_ranks:<2d} " + f"{num_recv_tokens:12d} " + f"{num_touched_experts:14d} | " + f"{compute_tflops:15.0f} " + f"{hbm_gbs_:9.0f} " + f"{nvlink_gbs_:9.0f} " + f"{t * 1e6:9.0f} " + f"{reduction} " + f"{speedup_text}" + ) + + dist_print("Performance:", once_in_node=True) + dist_print( + " > kind EP recv_tokens active_experts | " + "compute(TFLOPS) HBM(GB/s) NVL(GB/s) time(us) reduction(us) speedup", + once_in_node=True, + ) + dist_print( + fmt_perf_line( + "[fused]", + t_fused, + tflops * approx_factor, + hbm_gbs * approx_factor, + nvlink_gbs * approx_factor, + reduction_us=t_reduction * 1e6, + ) + ) + if ep_buffer is not None: + speedup = safe_div(t_baseline, t_fused) + dist_print( + fmt_perf_line( + "[baseline]", + t_baseline, + tflops_baseline, + hbm_gbs_baseline, + nvlink_gbs_baseline, + speedup=speedup, + ) + ) + else: + reason = ( + "disabled; pass --run-baseline or --check-output-diff to compare" + if not run_baseline_enabled + else "deep_ep unavailable" + ) + dist_print(f" > [baseline] ({reason})", once_in_node=True) + if ll_buffer is not None: + speedup_ll = safe_div(t_baseline_ll, t_fused) + dist_print( + fmt_perf_line( + "[ll_base]", + t_baseline_ll, + tflops_baseline_ll, + hbm_gbs_baseline_ll, + nvlink_gbs_baseline_ll, + speedup=speedup_ll, + ) + ) + elif run_ll_baseline_enabled: + dist_print(" > [ll_base] (deep_ep unavailable)", once_in_node=True) + + # ---- Cleanup ---- + dist.barrier() + sym_buffer.destroy() + if ep_buffer is not None: + ep_buffer.destroy() + if ll_buffer is not None: + ll_buffer.destroy() + dist.destroy_process_group() + + +# ============================================================================ +# Section 8: argparse + spawn. +# ============================================================================ + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="SM90 MegaMoE: fused (deep_gemm.fp8_mega_moe) vs DeepEP+grouped-FP8 baseline" + ) + + # Resources. + parser.add_argument( + "--ncu-profile-only", + action="store_true", + help="Run the fused SM90 kernel once for NCU/Nsight profiling", + ) + parser.add_argument( + "--fused-only-sweep", + action="store_true", + help="Run the fused-only token sweep benchmark mode", + ) + parser.add_argument( + "--accuracy", + action="store_true", + help="Run the layered SM90 accuracy suite instead of benchmark modes", + ) + parser.add_argument( + "--num-processes", type=int, default=8, help="Number of spawned processes, one per GPU" + ) + parser.add_argument( + "--local-rank-idx", + type=int, + default=None, + help="Local rank for single-process mode, useful for external launchers/NCU", + ) + + # Model shape. + # SM90 fused kernel requires intermediate_hidden <= 4096. + parser.add_argument("--num-max-tokens-per-rank", type=int, default=8192) + parser.add_argument( + "--num-tokens", + type=int, + default=0, + help="Actual per-rank token count; 0 means num-max-tokens-per-rank", + ) + parser.add_argument( + "--num-max-removed-tokens", + type=int, + default=0, + help="Max random token removals per rank when num-tokens is 0", + ) + parser.add_argument( + "--batches", + type=int, + nargs="+", + default=None, + help="Token counts for --fused-only-sweep (default: 1 2 4 8 16 32)", + ) + parser.add_argument("--hidden", type=int, default=7168) + parser.add_argument( + "--intermediate-hidden", + type=int, + default=3072, + help="Intermediate dimension, constrained to <= 4096 by SM90 l2_arrival_mask", + ) + parser.add_argument( + "--activation-clamp", + type=float, + default=10.0, + help="Clamp threshold for gate/up before SwiGLU; pass inf to disable", + ) + parser.add_argument("--num-experts", type=int, default=384) + parser.add_argument("--num-topk", type=int, default=6) + parser.add_argument( + "--masked-ratio", + type=float, + default=0.0, + help="Randomly mask some topk expert selections to test sparse routing edges", + ) + parser.add_argument( + "--fast-math", + type=int, + default=1, + help="Whether fused SwiGLU uses fast math (0/1)", + ) + + # Timing. + parser.add_argument( + "--num-bench-tests", + "--num-tests", + dest="num_bench_tests", + type=int, + default=30, + help="Number of bench_kineto iterations for the fused kernel", + ) + parser.add_argument( + "--num-warmup", type=int, default=5, help="baseline cuda events warmup" + ) + parser.add_argument( + "--num-repeat", type=int, default=20, help="Baseline CUDA event timing iterations" + ) + parser.add_argument( + "--l2-flush-gb", + type=float, + default=8.0, + help="Temporary write size used to flush L2 before baseline timing; 0 disables it", + ) + parser.add_argument( + "--run-baseline", + action="store_true", + help="Enable the DeepEP+grouped-FP8 baseline; disabled by default", + ) + parser.add_argument( + "--run-low-latency-baseline", + action="store_true", + help=( + "Enable the sglang low-latency baseline " + "(DeepEP low_latency_dispatch -> masked grouped FP8 GEMM -> masked " + "SwiGLU+FP8 quant -> masked FP8 GEMM -> low_latency_combine); " + "disabled by default" + ), + ) + parser.add_argument( + "--check-output-diff", + type=int, + default=0, + help="If nonzero, print fused vs per-128 baseline output differences", + ) + parser.add_argument( + "--dump-profile-traces", + type=str, + default="", + help="If nonempty, write one fused Chrome trace per rank to this directory", + ) + + # Accuracy mode. + parser.add_argument( + "--layers", + type=int, + nargs="+", + default=[1, 2, 3, 4], + help="Accuracy layers to run with --accuracy (1..5); default: 1 2 3 4", + ) + parser.add_argument( + "--num-correctness-tests", + type=int, + default=None, + help="Layer-5 random stress scenario count for --accuracy", + ) + parser.add_argument( + "--filter", + type=str, + default="", + help="Substring filter for --accuracy scenario names", + ) + parser.add_argument( + "--diff-tol", + type=float, + default=0.07, + help="calc_diff tolerance for --accuracy; default: 0.07", + ) + parser.add_argument( + "--fail-fast", + action="store_true", + help="Stop --accuracy on the first failing scenario", + ) + + args = parser.parse_args() + + if args.dump_profile_traces: + os.makedirs(args.dump_profile_traces, exist_ok=True) + + if args.local_rank_idx is not None: + # Single-process mode: external launcher sets MASTER_ADDR/PORT/WORLD_SIZE/RANK. + test(args.local_rank_idx, args.num_processes, args) + else: + # Multi-process mode: one process per GPU; test() creates the NCCL group. + torch.multiprocessing.spawn( + test, args=(args.num_processes, args), nprocs=args.num_processes + ) From a36f7fd6cf8da5f020980dd897b828c2cf8e59a9 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 15 Jun 2026 16:15:27 -0700 Subject: [PATCH 21/26] Move MegaMoE Hopper test into sgl_deep_gemm tests (#45) --- {tests => sgl_deep_gemm/tests}/test_mega_moe_hopper.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename {tests => sgl_deep_gemm/tests}/test_mega_moe_hopper.py (100%) diff --git a/tests/test_mega_moe_hopper.py b/sgl_deep_gemm/tests/test_mega_moe_hopper.py similarity index 100% rename from tests/test_mega_moe_hopper.py rename to sgl_deep_gemm/tests/test_mega_moe_hopper.py From b5238adf3ec6cd52882126fee97edc7a47c4e517 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 15 Jun 2026 16:16:33 -0700 Subject: [PATCH 22/26] Update test modification instruction Clarify pull request rebase instructions and testing location. --- sgl_deep_gemm/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/sgl_deep_gemm/README.md b/sgl_deep_gemm/README.md index 39d56d391e..1be1980b9c 100644 --- a/sgl_deep_gemm/README.md +++ b/sgl_deep_gemm/README.md @@ -13,4 +13,4 @@ To release a new set of wheels, please contact SGLang team and run the [release For each major version release (0.X.Y -> 0.(X+1).0), a new branch should be created (release/v0.(X+1).0) for stability purpose. -For any incoming pull requests, it should be rebased upon `dev` branch. +For any incoming pull requests, it should be rebased upon `dev` branch. Any newly added or modified tests should be put under `sgl_deep_gemm/tests` From bdabf6c3f18dae178c4c804027ce00b68ac94cf4 Mon Sep 17 00:00:00 2001 From: Baizhou Zhang Date: Mon, 15 Jun 2026 18:01:38 -0700 Subject: [PATCH 23/26] Add Hopper mega moe test to runner (#46) --- sgl_deep_gemm/run_tests.sh | 24 +++++++++++++++++++----- 1 file changed, 19 insertions(+), 5 deletions(-) diff --git a/sgl_deep_gemm/run_tests.sh b/sgl_deep_gemm/run_tests.sh index 20b5394a6b..d1c74b9ad8 100755 --- a/sgl_deep_gemm/run_tests.sh +++ b/sgl_deep_gemm/run_tests.sh @@ -37,7 +37,7 @@ if [ "${NUM_GPUS}" -eq 0 ]; then echo "No GPUs visible to nvidia-smi — DeepGEMM tests require a GPU." >&2 exit 1 fi -# arch major: 9 == Hopper (SM90), 10 == Blackwell (SM100). +# arch major: 9 == Hopper (SM90), 10 == Blackwell (SM100/SM103). COMPUTE_CAP=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null | head -1) ARCH_MAJOR=${COMPUTE_CAP%%.*} @@ -123,15 +123,23 @@ if [ -f "${TESTS_DIR}/test_lazy_init.py" ]; then skip_test test_lazy_init.py "tvm_ffi eager CUDA init (upstream)" fi -# mega_moe family uses SM100 fp4 + symmetric-memory kernels (SM100-only). +# mega_moe family covers Blackwell fp4 + symmetric-memory kernels and the SM90 +# Hopper fp8 path. # test_mega_moe.py additionally needs deep_ep (with ElasticBuffer); the l1 and # pre_dispatch tests use deep_gemm's own symmetric buffer. -MEGA_MOE_ALL=( +MEGA_MOE_BLACKWELL=( test_mega_moe.py test_mega_moe_l1_fp4_accuracy.py test_mega_moe_l1_sentinel.py test_mega_moe_pre_dispatch.py ) +MEGA_MOE_HOPPER=( + test_mega_moe_hopper.py +) +MEGA_MOE_ALL=( + "${MEGA_MOE_BLACKWELL[@]}" + "${MEGA_MOE_HOPPER[@]}" +) MEGA_MOE_L1=( test_mega_moe_l1_fp4_accuracy.py test_mega_moe_l1_sentinel.py @@ -156,9 +164,15 @@ elif [ "${ARCH_MAJOR}" -ge 10 ]; then [ -f "${TESTS_DIR}/${t}" ] && skip_test "${t}" "fp4 kernel failures, see comment (confirm on B200)" done [ -f "${TESTS_DIR}/test_mega_moe_pre_dispatch.py" ] && run_test test_mega_moe_pre_dispatch.py + for t in "${MEGA_MOE_HOPPER[@]}"; do + [ -f "${TESTS_DIR}/${t}" ] && skip_test "${t}" "SM90-only, arch major ${ARCH_MAJOR}" + done else - for t in "${MEGA_MOE_ALL[@]}"; do - [ -f "${TESTS_DIR}/${t}" ] && skip_test "${t}" "SM100-only, arch major ${ARCH_MAJOR}" + for t in "${MEGA_MOE_BLACKWELL[@]}"; do + [ -f "${TESTS_DIR}/${t}" ] && skip_test "${t}" "Blackwell-only, arch major ${ARCH_MAJOR}" + done + for t in "${MEGA_MOE_HOPPER[@]}"; do + [ -f "${TESTS_DIR}/${t}" ] && run_test "${t}" --num-processes "${NPROC}" done fi From 77c952222cc8b9c3028b41530daeeb802117c27c Mon Sep 17 00:00:00 2001 From: Martin Hua Date: Thu, 18 Jun 2026 13:17:15 -0700 Subject: [PATCH 24/26] chore: bump apache-tvm-ffi 0.1.9 -> 0.1.11 (#47) --- setup.py | 2 +- sgl_deep_gemm/pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index e89d85fd7f..56a016696a 100644 --- a/setup.py +++ b/setup.py @@ -220,6 +220,6 @@ def run(self): 'build_py': CustomBuildPy, }, install_requires=[ - 'apache-tvm-ffi==0.1.9', + 'apache-tvm-ffi==0.1.11', ], ) diff --git a/sgl_deep_gemm/pyproject.toml b/sgl_deep_gemm/pyproject.toml index 36d89a94a4..40ce1c4e75 100644 --- a/sgl_deep_gemm/pyproject.toml +++ b/sgl_deep_gemm/pyproject.toml @@ -19,7 +19,7 @@ classifiers = [ dynamic = ["version"] dependencies = [ - "apache-tvm-ffi==0.1.9", + "apache-tvm-ffi==0.1.11", ] [project.optional-dependencies] From 774d08160d25d3a03db9f7abbde86053d3d6aeea Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 22 Jun 2026 18:03:54 +0000 Subject: [PATCH 25/26] feat: add signal for SBO in SM90 masked gemm Port the Single Batch Overlap (SBO) signal mechanism from the archive-release branch onto the refactored dev base. - sm90_fp8_gemm_1d2d.cuh: kEnableOverlap template param + signal pointer; emit a per-(group, m-block) completion counter via ptx::atomic_add_rel after the TMA store lands (ptx::tma_store_wait<0>). - GemmDesc: max_block_n / enable_overlap; SM90 heuristic caps block_n. - sm90 masked launch + gemm.hpp API return (block_m, signal_threshold); wired through both the pybind (_C) and tvm-ffi bindings + Python wrappers. --- csrc/apis/gemm.hpp | 25 +++++++++++++++---- csrc/jit_kernels/heuristics/config.hpp | 6 +++++ csrc/jit_kernels/heuristics/sm90.hpp | 1 + csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp | 25 ++++++++++++++----- csrc/tvm_ffi_api.cpp | 15 ++++++++--- deep_gemm/__init__.py | 6 +++-- .../deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh | 15 +++++++++-- sgl_deep_gemm/__init__.py | 6 +++-- 8 files changed, 78 insertions(+), 21 deletions(-) diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 41f16a5c3a..0811796670 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -219,7 +219,7 @@ static void m_grouped_fp8_fp4_gemm_nn_contiguous(const std::pair& a, +static std::optional> m_grouped_fp8_fp4_gemm_nt_masked(const std::pair& a, const std::pair& b, const torch::Tensor& d, const torch::Tensor& masked_m, @@ -228,13 +228,22 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair> recipe_a, std::optional> recipe_b, const std::string& compiled_dims, - const bool& disable_ue8m0_cast) { + const bool& disable_ue8m0_cast, + const int& max_block_n, + const bool& enable_overlap, + const std::optional& signal) { // Shape must be `[G, M, K] @ [G, N, K].mT` const auto major_a = get_major_type_ab(a.first); const auto major_b = get_major_type_ab(b.first); DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); DG_HOST_ASSERT(masked_m.is_contiguous()); + if (enable_overlap) { + DG_HOST_ASSERT(signal.has_value()); + DG_HOST_ASSERT(signal.value().is_contiguous()); + DG_HOST_ASSERT(signal.value().scalar_type() == torch::kInt32); + } + // Type and shape checks const auto arch_major = device_runtime->get_arch_major(); const auto [num_groups , m , k ] = check_grouped_ab_fp8_fp4(a.first, major_a, arch_major); @@ -255,17 +264,21 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair> result = std::nullopt; if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) { const auto major_sfb = get_major_type_ab(sfb); - sm90_m_grouped_fp8_gemm_masked_1d2d(a.first, sfa, b.first, sfb, d, masked_m, - num_groups, m, n, k, expected_m, major_a, major_b, major_sfb, compiled_dims); + result = sm90_m_grouped_fp8_gemm_masked_1d2d(a.first, sfa, b.first, sfb, d, masked_m, + num_groups, m, n, k, expected_m, major_a, major_b, major_sfb, compiled_dims, + max_block_n, enable_overlap, signal); } else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) { + DG_HOST_ASSERT(not enable_overlap and "SBO overlap is only supported on SM90"); sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(a.first, sfa, b.first, sfb, d, masked_m, num_groups, m, n, k, expected_m, gran_k_a, gran_k_b, major_a, major_b, compiled_dims); } else { DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types"); } + return result; } static void k_grouped_fp8_gemm_tn_contiguous(const std::pair& a, @@ -645,7 +658,9 @@ static void register_apis(pybind11::module_& m) { py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"), py::arg("expected_m"), py::arg("recipe") = std::nullopt, py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, - py::arg("compiled_dims") = "nk", py::arg("disable_ue8m0_cast") = false); + py::arg("compiled_dims") = "nk", py::arg("disable_ue8m0_cast") = false, + py::arg("max_block_n") = 256, py::arg("enable_overlap") = false, + py::arg("signal") = std::nullopt); m.def("k_grouped_fp8_gemm_tn_contiguous", &k_grouped_fp8_gemm_tn_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("ks"), py::arg("ks_tensor"), py::arg("c") = std::nullopt, diff --git a/csrc/jit_kernels/heuristics/config.hpp b/csrc/jit_kernels/heuristics/config.hpp index c06f2f168f..3d307737b1 100644 --- a/csrc/jit_kernels/heuristics/config.hpp +++ b/csrc/jit_kernels/heuristics/config.hpp @@ -22,6 +22,10 @@ struct GemmDesc { int num_sms, tc_util; std::string compiled_dims; + // SBO (Single Batch Overlap), SM90 masked GEMM only + int max_block_n = 256; + bool enable_overlap = false; + // Shape for heuristic generation int expected_m = 0, expected_n = 0, expected_k = 0, expected_num_groups = 0; int get_expected_m() const { return expected_m > 0 ? expected_m : m; } @@ -60,6 +64,8 @@ struct GemmDesc { << ", num_sms=" << desc.num_sms << ", tc_util=" << desc.tc_util << ", compiled_dims=" << desc.compiled_dims + << ", max_block_n=" << desc.max_block_n + << ", enable_overlap=" << static_cast(desc.enable_overlap) << ", expected_m=" << desc.expected_m << ", expected_n=" << desc.expected_n << ", expected_k=" << desc.expected_k diff --git a/csrc/jit_kernels/heuristics/sm90.hpp b/csrc/jit_kernels/heuristics/sm90.hpp index c411fb7e01..e7ac205869 100644 --- a/csrc/jit_kernels/heuristics/sm90.hpp +++ b/csrc/jit_kernels/heuristics/sm90.hpp @@ -52,6 +52,7 @@ struct SM90ArchSpec { end = 192; if (desc.kernel_type == KernelType::Kernel1D1D) end = 160; + end = std::min(end, desc.max_block_n); // Enumerate for (int i = start; i <= end; i += step) block_n_candidates.push_back(i); diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index 61b85ec1f0..f8b35edf08 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -24,7 +24,7 @@ class SM90FP8Gemm1D2DRuntime final: public LaunchRuntime const std::optional& epilogue_type; cute::UMMA::Major major_sfb; - void *sfb, *grouped_layout; + void *sfb, *grouped_layout, *signal; CUtensorMap tensor_map_a; CUtensorMap tensor_map_b; CUtensorMap tensor_map_d; @@ -48,7 +48,7 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, {}, - {} + {}, {} >); }}; )", @@ -64,13 +64,14 @@ static void __instantiate_kernel() {{ args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads, args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type), - get_default_epilogue_type(args.epilogue_type)); + get_default_epilogue_type(args.epilogue_type), + args.gemm_desc.enable_overlap); } static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { // TODO: optimize `args` copy DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.sfb, args.grouped_layout, + args.sfb, args.grouped_layout, args.signal, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.tensor_map_sfa)); @@ -133,6 +134,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, .major_sfb = major_sfb, .sfb = sfb.data_ptr(), .grouped_layout = nullptr, + .signal = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, @@ -210,6 +212,7 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons .major_sfb = major_sfb, .sfb = sfb.data_ptr(), .grouped_layout = m_indices.data_ptr(), + .signal = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, @@ -220,14 +223,17 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons SM90FP8Gemm1D2DRuntime::launch(runtime, args); } -static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, +static std::optional> sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, const torch::Tensor& b, const torch::Tensor& sfb, const torch::Tensor& d, const torch::Tensor& masked_m, const int& num_groups, const int& m, const int& n, const int& k, const int& expected_m, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb, - const std::string& compiled_dims) { + const std::string& compiled_dims, + const int& max_block_n, + const bool& enable_overlap, + const std::optional& signal) { DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); @@ -241,6 +247,7 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to .with_accumulation = false, .num_sms = device_runtime->get_num_sms(), .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .max_block_n = max_block_n, .enable_overlap = enable_overlap, .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups }; const auto config = get_best_config(desc); @@ -277,6 +284,7 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to .major_sfb = major_sfb, .sfb = sfb.data_ptr(), .grouped_layout = masked_m.data_ptr(), + .signal = enable_overlap ? signal.value().data_ptr() : nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, @@ -285,6 +293,10 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to const auto code = SM90FP8Gemm1D2DRuntime::generate(args); const auto runtime = compiler->build("sm90_fp8_m_grouped_gemm_masked_1d2d", code); SM90FP8Gemm1D2DRuntime::launch(runtime, args); + + return enable_overlap + ? std::optional(std::make_pair(config.layout.block_m, ceil_div(n, config.layout.block_n))) + : std::nullopt; } static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, @@ -348,6 +360,7 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, .major_sfb = major_sfb, .sfb = sfb.data_ptr(), .grouped_layout = nullptr, + .signal = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, diff --git a/csrc/tvm_ffi_api.cpp b/csrc/tvm_ffi_api.cpp index 4031aec1bd..3ef636d741 100644 --- a/csrc/tvm_ffi_api.cpp +++ b/csrc/tvm_ffi_api.cpp @@ -336,7 +336,7 @@ void dg_m_grouped_fp8_fp4_gemm_nn_contiguous(TensorView a, TensorView a_sf, ); } -void dg_m_grouped_fp8_fp4_gemm_nt_masked(TensorView a, TensorView a_sf, +Tuple, Optional> dg_m_grouped_fp8_fp4_gemm_nt_masked(TensorView a, TensorView a_sf, TensorView b, TensorView b_sf, TensorView d, TensorView masked_m, @@ -345,17 +345,24 @@ void dg_m_grouped_fp8_fp4_gemm_nt_masked(TensorView a, TensorView a_sf, Optional> recipe_a, Optional> recipe_b, std::string compiled_dims, - bool disable_ue8m0_cast) { + bool disable_ue8m0_cast, + int64_t max_block_n, + bool enable_overlap, + Optional signal) { auto recipe_opt = recipe.has_value()? std::make_optional(std::make_tuple((int)recipe.value().get<0>(), (int)recipe.value().get<1>(), (int)recipe.value().get<2>())) : std::nullopt; auto recipe_a_opt = recipe_a.has_value() ? std::make_optional(std::make_tuple((int)recipe_a.value().get<0>(), (int)recipe_a.value().get<1>())) : std::nullopt; auto recipe_b_opt = recipe_b.has_value() ? std::make_optional(std::make_tuple((int)recipe_b.value().get<0>(), (int)recipe_b.value().get<1>())) : std::nullopt; - gemm::m_grouped_fp8_fp4_gemm_nt_masked( + auto signal_opt = signal.has_value() ? std::optional(convert_to_torch_tensor(signal.value())) : std::nullopt; + auto result = gemm::m_grouped_fp8_fp4_gemm_nt_masked( std::make_pair(convert_to_torch_tensor(a), convert_to_torch_tensor(a_sf)), std::make_pair(convert_to_torch_tensor(b), convert_to_torch_tensor(b_sf)), convert_to_torch_tensor(d), convert_to_torch_tensor(masked_m), (int) expected_m, recipe_opt, recipe_a_opt, recipe_b_opt, - compiled_dims, disable_ue8m0_cast + compiled_dims, disable_ue8m0_cast, (int) max_block_n, enable_overlap, signal_opt ); + if (result.has_value()) + return Tuple, Optional>(Optional(result->first), Optional(result->second)); + return Tuple, Optional>(Optional(), Optional()); } diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index 2bce4038da..848982e67d 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -249,9 +249,11 @@ def transform_sf_into_required_layout(sf, mn, k, recipe, num_groups=None, is_sfa get_mk_alignment_for_contiguous_layout = _C.get_mk_alignment_for_contiguous_layout - def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False): + def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False, + max_block_n=256, enable_overlap=False, signal=None): (a, a_sf), (b, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) - return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, + max_block_n, enable_overlap, signal) fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_fp4_gemm_nt_masked diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh index aa412484de..b561f5e5db 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh @@ -43,9 +43,9 @@ template + typename epilogue_type_t, bool kEnableOverlap> CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void -sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, +sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, int* signal, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, const __grid_constant__ cute::TmaDescriptor tensor_map_a, const __grid_constant__ cute::TmaDescriptor tensor_map_b, @@ -436,6 +436,17 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, cute::tma_store_arrive(); } __syncwarp(); + + if constexpr (kEnableOverlap) { + if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) + cute::tma_store_wait<0>(); + + cutlass::arch::NamedBarrier::sync(kNumWGMMAStoreThreads, 1); + + if (threadIdx.x == 0) + ptx::atomic_add_rel(reinterpret_cast( + signal + scheduler.current_group_idx * math::ceil_div(shape_m, BLOCK_M) + m_block_idx), 1u); + } } } #else diff --git a/sgl_deep_gemm/__init__.py b/sgl_deep_gemm/__init__.py index 21bf6a7ed1..76ff2d1595 100644 --- a/sgl_deep_gemm/__init__.py +++ b/sgl_deep_gemm/__init__.py @@ -246,9 +246,11 @@ def transform_sf_into_required_layout(sf, mn, k, recipe, num_groups=None, is_sfa get_mk_alignment_for_contiguous_layout = _C.get_mk_alignment_for_contiguous_layout - def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False): + def m_grouped_fp8_fp4_gemm_nt_masked(a, b, d, masked_m, expected_m, recipe=None, recipe_a=None, recipe_b=None, compiled_dims='nk', disable_ue8m0_cast=False, + max_block_n=256, enable_overlap=False, signal=None): (a, a_sf), (b, b_sf) = _parse_tensor_or_tuple(a), _parse_tensor_or_tuple(b) - return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast) + return _C.m_grouped_fp8_fp4_gemm_nt_masked(a, a_sf, b, b_sf, d, masked_m, expected_m, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, + max_block_n, enable_overlap, signal) fp8_m_grouped_gemm_nt_masked = m_grouped_fp8_fp4_gemm_nt_masked From f4945a5ddff49e13d21912247f82910a1d042905 Mon Sep 17 00:00:00 2001 From: Brayden Zhong Date: Mon, 22 Jun 2026 18:03:54 +0000 Subject: [PATCH 26/26] feat: add test for signal GEMM SM90-guarded overlap test for m-grouped masked GEMM plus check_signal helper; mirrored into both the source and sgl_deep_gemm test trees. --- deep_gemm/testing/numeric.py | 16 +++++++++++ sgl_deep_gemm/tests/test_fp8_fp4.py | 41 +++++++++++++++++++++++++++-- tests/test_fp8_fp4.py | 41 +++++++++++++++++++++++++++-- 3 files changed, 94 insertions(+), 4 deletions(-) diff --git a/deep_gemm/testing/numeric.py b/deep_gemm/testing/numeric.py index a42c4318db..f98c3f12fd 100644 --- a/deep_gemm/testing/numeric.py +++ b/deep_gemm/testing/numeric.py @@ -19,3 +19,19 @@ def count_bytes(*tensors): elif t is not None: total += t.numel() * t.element_size() return total + + +def check_signal(num_local_expert, max_m, block_m, threshold, signal, masked_m): + ceil_div = lambda a, b: (a + b - 1) // b + + expert_len = max_m // block_m + for expert in range(num_local_expert): + mask = masked_m[expert] + start = expert * expert_len + end = expert * expert_len + expert_len + valid_len = ceil_div(mask, block_m) + for i in range(start, end): + if i < start + valid_len: + assert signal[i] == threshold, f'{i=}, {signal[i]=}, {threshold=}' + else: + assert signal[i] == 0, f'{i=}, {signal[i]=}' diff --git a/sgl_deep_gemm/tests/test_fp8_fp4.py b/sgl_deep_gemm/tests/test_fp8_fp4.py index 42038e459d..0906998e8c 100644 --- a/sgl_deep_gemm/tests/test_fp8_fp4.py +++ b/sgl_deep_gemm/tests/test_fp8_fp4.py @@ -6,12 +6,12 @@ import deep_gemm from deep_gemm.testing import ( bench_kineto, - calc_diff, count_bytes, + calc_diff, count_bytes, check_signal, ignore_env, get_arch_major ) from generators import ( - KernelType, get_ue8m0_usage, layout_masked_to_psum, align, + KernelType, QuantConfig, get_ue8m0_usage, layout_masked_to_psum, align, enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous, generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous, get_mk_alignment_for_contiguous_layout @@ -174,6 +174,42 @@ def test_func(): print() +def test_m_grouped_gemm_masked_overlap() -> None: + if get_arch_major() != 9: + print('Skipping m-grouped masked GEMM SBO overlap test (SM90 only)\n') + return + + print('Testing m-grouped masked GEMM with SBO overlap:') + ceil_div = lambda a, b: (a + b - 1) // b + + max_m = 4096 + quant_config = QuantConfig() + use_ue8m0 = get_ue8m0_usage(KernelType.Kernel1D2D) + disable_ue8m0_cast = not use_ue8m0 + recipe, recipe_a, recipe_b = quant_config.get_recipes() + + for num_groups, expected_m_per_group in ((1, 1024), (2, 512), (4, 256), (16, 64), (16, 32)): + for n, k in ((4096, 7168), (7168, 2048)): + for i in range(10): + a, b, masked_m, _, d, ref_d = generate_m_grouped_masked( + num_groups, max_m, expected_m_per_group, n, k, + use_ue8m0=use_ue8m0, quant_config=quant_config) + + signal = torch.zeros(num_groups * ceil_div(max_m, 64), dtype=torch.int32, device='cuda') + block_m, threshold = deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked( + a, b, d, masked_m, int(expected_m_per_group * 1.2), + disable_ue8m0_cast=disable_ue8m0_cast, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b, + enable_overlap=True, signal=signal) + + check_signal(num_groups, max_m, block_m, threshold, signal, masked_m) + for j in range(num_groups): + diff = calc_diff(d[j, :masked_m[j].item()], ref_d[j, :masked_m[j].item()]) + assert diff < quant_config.max_diff(), f'{n=}, {k=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}' + print(f' > Passed (num_groups={num_groups:2}, expected_m_per_group={expected_m_per_group:4})') + print() + + def test_k_grouped_gemm_contiguous() -> None: print('Testing k-grouped contiguous GEMM:') @@ -220,4 +256,5 @@ def test_func(): # test_gemm() test_m_grouped_gemm_contiguous() # test_m_grouped_gemm_masked() + test_m_grouped_gemm_masked_overlap() # test_k_grouped_gemm_contiguous() diff --git a/tests/test_fp8_fp4.py b/tests/test_fp8_fp4.py index 4e9f54f7f4..78142f10e8 100644 --- a/tests/test_fp8_fp4.py +++ b/tests/test_fp8_fp4.py @@ -6,12 +6,12 @@ import deep_gemm from deep_gemm.testing import ( bench_kineto, - calc_diff, count_bytes, + calc_diff, count_bytes, check_signal, ignore_env, get_arch_major ) from generators import ( - KernelType, get_ue8m0_usage, layout_masked_to_psum, align, + KernelType, QuantConfig, get_ue8m0_usage, layout_masked_to_psum, align, enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous, generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous, get_mk_alignment_for_contiguous_layout @@ -173,6 +173,42 @@ def test_func(): print() +def test_m_grouped_gemm_masked_overlap() -> None: + if get_arch_major() != 9: + print('Skipping m-grouped masked GEMM SBO overlap test (SM90 only)\n') + return + + print('Testing m-grouped masked GEMM with SBO overlap:') + ceil_div = lambda a, b: (a + b - 1) // b + + max_m = 4096 + quant_config = QuantConfig() + use_ue8m0 = get_ue8m0_usage(KernelType.Kernel1D2D) + disable_ue8m0_cast = not use_ue8m0 + recipe, recipe_a, recipe_b = quant_config.get_recipes() + + for num_groups, expected_m_per_group in ((1, 1024), (2, 512), (4, 256), (16, 64), (16, 32)): + for n, k in ((4096, 7168), (7168, 2048)): + for i in range(10): + a, b, masked_m, _, d, ref_d = generate_m_grouped_masked( + num_groups, max_m, expected_m_per_group, n, k, + use_ue8m0=use_ue8m0, quant_config=quant_config) + + signal = torch.zeros(num_groups * ceil_div(max_m, 64), dtype=torch.int32, device='cuda') + block_m, threshold = deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked( + a, b, d, masked_m, int(expected_m_per_group * 1.2), + disable_ue8m0_cast=disable_ue8m0_cast, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b, + enable_overlap=True, signal=signal) + + check_signal(num_groups, max_m, block_m, threshold, signal, masked_m) + for j in range(num_groups): + diff = calc_diff(d[j, :masked_m[j].item()], ref_d[j, :masked_m[j].item()]) + assert diff < quant_config.max_diff(), f'{n=}, {k=}, {j=}, masked_m={masked_m[j]}, {num_groups=}, {diff:.5f}' + print(f' > Passed (num_groups={num_groups:2}, expected_m_per_group={expected_m_per_group:4})') + print() + + def test_k_grouped_gemm_contiguous() -> None: print('Testing k-grouped contiguous GEMM:') @@ -219,4 +255,5 @@ def test_func(): test_gemm() test_m_grouped_gemm_contiguous() test_m_grouped_gemm_masked() + test_m_grouped_gemm_masked_overlap() test_k_grouped_gemm_contiguous()