From 66081d4c9c7d7c44f13fea402e5b622aa0f409c2 Mon Sep 17 00:00:00 2001 From: Zhean Xu <94977922+zheanxu@users.noreply.github.com> Date: Thu, 10 Sep 2026 14:31:04 +0800 Subject: [PATCH 1/5] Public Release 26/09 (#432) * Public release 26/09 * Update News --- .gitmodules | 6 +- CMakeLists.txt | 6 +- README.md | 42 +- csrc/apis/attention.hpp | 288 +++++-- csrc/apis/{runtime.hpp => config.hpp} | 38 +- csrc/apis/einsum.hpp | 119 ++- csrc/apis/gemm.hpp | 300 ++++++-- csrc/apis/hyperconnection.hpp | 8 +- csrc/apis/layout.hpp | 21 +- csrc/apis/mega_gate.hpp | 234 ++++++ csrc/apis/mega_mhc.hpp | 287 +++++++ csrc/apis/{mega.hpp => mega_moe.hpp} | 25 +- csrc/indexing/main.cu | 2 +- csrc/jit/cache.hpp | 31 - csrc/jit/compiler.hpp | 362 --------- csrc/jit/device_runtime.hpp | 138 ---- csrc/jit/handle.hpp | 222 ------ csrc/jit/include_parser.hpp | 80 -- csrc/jit/kernel_runtime.hpp | 165 ---- csrc/jit_kernels/heuristics/common.hpp | 8 +- csrc/jit_kernels/heuristics/config.hpp | 21 +- csrc/jit_kernels/heuristics/mega_gate.hpp | 236 ++++++ csrc/jit_kernels/heuristics/mega_moe.hpp | 66 +- csrc/jit_kernels/heuristics/runtime.hpp | 29 +- csrc/jit_kernels/heuristics/sm100.hpp | 75 +- csrc/jit_kernels/heuristics/sm90.hpp | 5 +- csrc/jit_kernels/heuristics/utils.hpp | 9 + csrc/jit_kernels/impls/epilogue.hpp | 33 + csrc/jit_kernels/impls/runtime_utils.hpp | 49 +- csrc/jit_kernels/impls/sm100_bf16_gemm.hpp | 189 ++--- .../impls/sm100_bf16_mega_gate.hpp | 123 +++ .../jit_kernels/impls/sm100_bf16_mega_moe.hpp | 195 ++--- csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp | 105 +-- .../impls/sm100_fp8_fp4_gemm_1d1d.hpp | 313 +++++--- .../impls/sm100_fp8_fp4_mega_moe.hpp | 239 ++---- csrc/jit_kernels/impls/sm100_mega_mhc.hpp | 149 ++++ csrc/jit_kernels/impls/sm100_mqa_logits.hpp | 592 +++++++-------- .../impls/sm100_sparse_mqa_logits.hpp | 202 +++++ .../impls/sm100_tf32_hc_prenorm_gemm.hpp | 108 +-- csrc/jit_kernels/impls/sm90_bf16_gemm.hpp | 164 ++-- csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp | 103 +-- csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp | 73 +- csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp | 118 +-- .../jit_kernels/impls/sm90_fp8_mqa_logits.hpp | 372 +++------ .../impls/sm90_tf32_hc_prenorm_gemm.hpp | 108 +-- csrc/jit_kernels/impls/smxx_clean_logits.hpp | 81 -- csrc/jit_kernels/impls/smxx_cublaslt.hpp | 118 ++- csrc/jit_kernels/impls/smxx_layout.hpp | 228 +++--- csrc/python_api.cpp | 19 +- csrc/runtime/jit.hpp | 34 + csrc/runtime/runtime.hpp | 99 +++ csrc/utils/compatibility.hpp | 15 +- csrc/utils/exception.hpp | 48 +- csrc/utils/format.hpp | 6 - csrc/utils/hash.hpp | 39 - csrc/utils/layout.hpp | 22 +- csrc/utils/lazy_init.hpp | 27 - csrc/utils/system.hpp | 128 ---- deep_gemm/__init__.py | 132 ++-- deep_gemm/include/deep_gemm/comm/barrier.cuh | 37 +- deep_gemm/include/deep_gemm/common/math.cuh | 135 +++- .../include/deep_gemm/common/packing.cuh | 14 + .../deep_gemm/common/ring_pipeline.cuh | 40 + .../include/deep_gemm/common/tma_copy.cuh | 58 +- deep_gemm/include/deep_gemm/common/types.cuh | 19 +- deep_gemm/include/deep_gemm/common/utils.cuh | 25 - .../deep_gemm/epilogue/clean_logits.cuh | 49 ++ .../deep_gemm/epilogue/sm100_mega_gate.cuh | 282 +++++++ .../deep_gemm/epilogue/sm100_mega_mhc.cuh | 428 +++++++++++ .../deep_gemm/epilogue/sm100_store_cd.cuh | 172 +++-- .../epilogue/sm100_store_cd_swap_ab.cuh | 142 +++- .../include/deep_gemm/epilogue/transform.cuh | 53 +- .../deep_gemm/impls/sm100_bf16_gemm.cuh | 176 ++--- .../deep_gemm/impls/sm100_bf16_mega_gate.cuh | 483 ++++++++++++ .../deep_gemm/impls/sm100_bf16_mega_moe.cuh | 257 ++++--- .../impls/sm100_fp8_fp4_gemm_1d1d.cuh | 397 +++++----- .../impls/sm100_fp8_fp4_mega_moe.cuh | 300 ++++---- .../deep_gemm/impls/sm100_mega_mhc.cuh | 704 ++++++++++++++++++ .../deep_gemm/impls/sm100_mqa_logits.cuh | 172 +++-- .../impls/sm100_sparse_mqa_logits.cuh | 627 ++++++++++++++++ .../impls/sm100_tf32_hc_prenorm_gemm.cuh | 1 + .../deep_gemm/impls/sm90_bf16_gemm.cuh | 6 + .../deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh | 10 +- .../deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh | 4 +- .../deep_gemm/impls/sm90_fp8_mqa_logits.cuh | 20 +- .../deep_gemm/impls/smxx_clean_logits.cuh | 74 -- .../include/deep_gemm/impls/smxx_layout.cuh | 133 ++-- deep_gemm/include/deep_gemm/layout/gemm.cuh | 67 ++ .../include/deep_gemm/layout/mega_gate.cuh | 130 ++++ .../include/deep_gemm/layout/mega_mhc.cuh | 169 +++++ .../include/deep_gemm/layout/mega_moe.cuh | 185 ++--- .../include/deep_gemm/layout/mqa_logits.cuh | 17 +- .../deep_gemm/layout/sparse_mqa_logits.cuh | 146 ++++ deep_gemm/include/deep_gemm/mma/sm100.cuh | 23 +- deep_gemm/include/deep_gemm/ptx/ld_st.cuh | 68 +- deep_gemm/include/deep_gemm/ptx/tcgen05.cuh | 51 ++ deep_gemm/include/deep_gemm/ptx/tma.cuh | 47 +- deep_gemm/include/deep_gemm/ptx/utils.cuh | 24 + .../include/deep_gemm/scheduler/gemm.cuh | 83 +-- .../include/deep_gemm/scheduler/mega_gate.cuh | 55 ++ .../include/deep_gemm/scheduler/mega_mhc.cuh | 55 ++ .../include/deep_gemm/scheduler/mega_moe.cuh | 40 +- .../deep_gemm/scheduler/sm100_mqa_logits.cuh | 138 +++- .../scheduler/sm100_mqa_logits_metadata.cuh | 121 +++ .../scheduler/sm100_paged_mqa_logits.cuh | 211 ++++-- .../sm100_sparse_mqa_logits_metadata.cuh | 468 ++++++++++++ deep_gemm/legacy/m_grouped_gemm.py | 2 +- deep_gemm/mega/__init__.py | 47 +- deep_gemm/utils/layout.py | 16 +- deep_gemm/utils/math.py | 7 +- docs/scaling-factor-format.md | 469 ++++++++++++ setup.py | 20 +- tests/generators.py | 374 ++++++---- tests/test_attention.py | 380 +++++++++- tests/test_bf16.py | 255 ++++--- tests/test_einsum.py | 93 ++- tests/test_fp8_fp4.py | 352 ++++++--- tests/test_layout.py | 66 +- tests/test_legacy.py | 9 +- tests/test_mega_gate.py | 354 +++++++++ tests/test_mega_mhc.py | 266 +++++++ tests/test_mega_moe.py | 48 +- tests/test_sanitizer.py | 2 +- tests/utils.py | 67 +- third-party/deep_jit | 1 + third-party/fmt | 1 - third-party/tilelang_ops/norm.py | 213 ++++++ third-party/tilelang_ops/ref_mhc.py | 277 +++++++ 128 files changed, 12176 insertions(+), 5113 deletions(-) rename csrc/apis/{runtime.hpp => config.hpp} (51%) create mode 100644 csrc/apis/mega_gate.hpp create mode 100644 csrc/apis/mega_mhc.hpp rename csrc/apis/{mega.hpp => mega_moe.hpp} (97%) delete mode 100644 csrc/jit/cache.hpp delete mode 100644 csrc/jit/compiler.hpp delete mode 100644 csrc/jit/device_runtime.hpp delete mode 100644 csrc/jit/handle.hpp delete mode 100644 csrc/jit/include_parser.hpp delete mode 100644 csrc/jit/kernel_runtime.hpp create mode 100644 csrc/jit_kernels/heuristics/mega_gate.hpp create mode 100644 csrc/jit_kernels/impls/sm100_bf16_mega_gate.hpp create mode 100644 csrc/jit_kernels/impls/sm100_mega_mhc.hpp create mode 100644 csrc/jit_kernels/impls/sm100_sparse_mqa_logits.hpp delete mode 100644 csrc/jit_kernels/impls/smxx_clean_logits.hpp create mode 100644 csrc/runtime/jit.hpp create mode 100644 csrc/runtime/runtime.hpp delete mode 100644 csrc/utils/format.hpp delete mode 100644 csrc/utils/hash.hpp delete mode 100644 csrc/utils/lazy_init.hpp delete mode 100644 csrc/utils/system.hpp create mode 100644 deep_gemm/include/deep_gemm/common/packing.cuh create mode 100644 deep_gemm/include/deep_gemm/common/ring_pipeline.cuh create mode 100644 deep_gemm/include/deep_gemm/epilogue/clean_logits.cuh create mode 100644 deep_gemm/include/deep_gemm/epilogue/sm100_mega_gate.cuh create mode 100644 deep_gemm/include/deep_gemm/epilogue/sm100_mega_mhc.cuh create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_gate.cuh create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_mega_mhc.cuh create mode 100644 deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh delete mode 100644 deep_gemm/include/deep_gemm/impls/smxx_clean_logits.cuh create mode 100644 deep_gemm/include/deep_gemm/layout/gemm.cuh create mode 100644 deep_gemm/include/deep_gemm/layout/mega_gate.cuh create mode 100644 deep_gemm/include/deep_gemm/layout/mega_mhc.cuh create mode 100644 deep_gemm/include/deep_gemm/layout/sparse_mqa_logits.cuh create mode 100644 deep_gemm/include/deep_gemm/scheduler/mega_gate.cuh create mode 100644 deep_gemm/include/deep_gemm/scheduler/mega_mhc.cuh create mode 100644 deep_gemm/include/deep_gemm/scheduler/sm100_mqa_logits_metadata.cuh create mode 100644 deep_gemm/include/deep_gemm/scheduler/sm100_sparse_mqa_logits_metadata.cuh create mode 100644 docs/scaling-factor-format.md create mode 100644 tests/test_mega_gate.py create mode 100644 tests/test_mega_mhc.py create mode 160000 third-party/deep_jit delete mode 160000 third-party/fmt create mode 100644 third-party/tilelang_ops/norm.py create mode 100644 third-party/tilelang_ops/ref_mhc.py diff --git a/.gitmodules b/.gitmodules index 332be63983..8921fa6ce1 100644 --- a/.gitmodules +++ b/.gitmodules @@ -1,6 +1,6 @@ [submodule "third-party/cutlass"] path = third-party/cutlass url = https://github.com/NVIDIA/cutlass.git -[submodule "third-party/fmt"] - path = third-party/fmt - url = https://github.com/fmtlib/fmt.git +[submodule "third-party/deep_jit"] + path = third-party/deep_jit + url = https://github.com/deepseek-ai/DeepJIT.git diff --git a/CMakeLists.txt b/CMakeLists.txt index bbf625d306..70797aa793 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1,5 +1,5 @@ # NOTES: current just for CMake-based IDE (e.g. CLion) indexing, the real compilation is done via JIT -cmake_minimum_required(VERSION 3.10) +cmake_minimum_required(VERSION 3.18) project(deep_gemm LANGUAGES CXX CUDA) set(CMAKE_VERBOSE_MAKEFILE ON) @@ -18,9 +18,11 @@ find_package(pybind11 REQUIRED) find_package(Torch REQUIRED) set(CMAKE_CXX_STANDARD 20) +set(CMAKE_CXX_STANDARD_REQUIRED ON) set(CMAKE_CUDA_STANDARD 20) +set(CMAKE_CUDA_STANDARD_REQUIRED ON) -include_directories(deep_gemm/include third-party/cutlass/include third-party/cutlass/tools/util/include third-party/fmt/include) +include_directories(deep_gemm/include third-party/deep_jit/include third-party/cutlass/include third-party/cutlass/tools/util/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) diff --git a/README.md b/README.md index 6ef705ffce..6d4e5b7989 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # DeepGEMM -DeepGEMM is a unified, high-performance tensor core kernel library that brings together the key computation primitives of modern large language models — GEMMs (FP8, FP4, BF16), fused MoE with overlapped communication (Mega MoE), MQA scoring for the lightning indexer, HyperConnection (HC), and more — into a single, cohesive CUDA codebase. All kernels are compiled at runtime via a lightweight Just-In-Time (JIT) module, requiring no CUDA compilation during installation. +DeepGEMM is a unified, high-performance tensor core kernel library that brings together the key computation primitives of modern large language models — GEMMs (FP8, FP4, BF16), fused MoE with overlapped communication (Mega MoE), MQA scoring for the lightning indexer, HyperConnection (HC), and more — into a single, cohesive CUDA codebase. All kernels are compiled at runtime through DeepJIT, requiring no CUDA compilation during installation. DeepGEMM leverages some concepts from [CUTLASS](https://github.com/nvidia/cutlass) and [CuTe](https://github.com/NVIDIA/cutlass/tree/main/include/cute), but avoids heavy reliance on their templates or algebras. The library is designed for simplicity, with only a limited number of core kernel functions, making it a clean and accessible resource for learning NVIDIA GPU kernel optimization techniques. @@ -8,18 +8,17 @@ Despite its lightweight design, DeepGEMM's performance matches or exceeds expert ## News +- 2026.09.10: Sparse Indexer, Mega Gate, Mega mHC, DeepJIT, MoE and Indexer optimizations and more. + - Please see [#432](https://github.com/deepseek-ai/DeepGEMM/pull/432) for more details. - 2026.04.16: Mega MoE, FP8xFP4 GEMM, FP4 Indexer, PDL, faster JIT compilation and more. - Please see [#304](https://github.com/deepseek-ai/DeepGEMM/pull/304) for more details. - For Mega MoE benchmarks, refer to [#316](https://github.com/deepseek-ai/DeepGEMM/pull/316). - 2025.09.28: DeepGEMM now supports scoring kernels (weighted ReLU MQA logits) for the lightning indexer for DeepSeek v3.2. - Please see [#200](https://github.com/deepseek-ai/DeepGEMM/pull/200) for more details. - 2025.07.20: DeepGEMM now supports both SM90/SM100, and has a full refactor with a low-CPU-overhead JIT CPP module. - - NVRTC and post-compilation SASS optimization are all disabled. - - NVRTC will be supported later. - As NVCC 12.9 will automatically do the FFMA interleaving, all post optimizations will be no longer supported. - Please see [#112](https://github.com/deepseek-ai/DeepGEMM/pull/112) for more details. - 2025.05.14: DeepGEMM now offers weight gradient kernels for dense and MoE backward! See [#95](https://github.com/deepseek-ai/DeepGEMM/pull/95) for details. -- 2025.05.07: DeepGEMM now supports NVRTC with up to 10x compilation speedup! See [#94](https://github.com/deepseek-ai/DeepGEMM/pull/94) for details. Please use `DG_JIT_USE_NVRTC=1` to enable it (may have performance loss with some cases). - 2025.04.18: DeepGEMM now achieves up to **1550 TFLOPS** on H800! See [#74](https://github.com/deepseek-ai/DeepGEMM/pull/74), [#78](https://github.com/deepseek-ai/DeepGEMM/pull/78), [#81](https://github.com/deepseek-ai/DeepGEMM/pull/81), [#86](https://github.com/deepseek-ai/DeepGEMM/pull/86) and [340d988](https://github.com/deepseek-ai/DeepGEMM/commit/340d9880f4a418d943d34260d20a79f41f4c0526) for details. ## Quick start @@ -28,14 +27,10 @@ Despite its lightweight design, DeepGEMM's performance matches or exceeds expert - NVIDIA SM90 or SM100 architecture GPU - Python 3.8 or higher -- Compilers with C++20 support -- CUDA Toolkit: - - CUDA 12.3 or higher for SM90 - - **We highly recommend 12.9 or higher for the best performance** - - CUDA 12.9 or higher for SM100 -- PyTorch 2.1 or higher +- Compilers and standard libraries with C++20 `` support +- CUDA Toolkit 12.9 or higher +- PyTorch 2.3 or higher - CUTLASS 4.0 or higher (could be cloned by Git submodule) -- `{fmt}` library (could be cloned by Git submodule) ### Development @@ -44,7 +39,7 @@ Despite its lightweight design, DeepGEMM's performance matches or exceeds expert git clone --recursive git@github.com:deepseek-ai/DeepGEMM.git cd DeepGEMM -# Link some essential includes and build the CPP JIT module +# Link some essential includes and build the C++ extension cat develop.sh ./develop.sh ``` @@ -113,16 +108,17 @@ For more details and the paged version `fp8_paged_mqa_logits`, please refer to ` #### Mega MoE -Mega MoE fuses and overlaps EP dispatch, linear 1 (FP8xFP4), SwiGLU, linear 2 (FP8xFP4), and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory. Usage: +Mega MoE fuses and overlaps EP dispatch, linear 1 and linear 2 (FP8xFP4 or FP8xFP8), SwiGLU, and EP combine into a single mega-kernel, overlapping NVLink communication and tensor core computation. It requires multi-process launch with symmetric memory. Usage: ```python # Allocate symmetric memory buffer # NOTES: requires PyTorch >= 2.9 buffer = deep_gemm.get_symm_buffer_for_mega_moe( - group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden + group, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, + mma_type='fp8xfp4', # Use 'fp8xfp8' for FP8 routed-expert weights ) -# Transform weights (FP4 with UE8M0 SF) into the required layout +# Transform routed weights (FP4 or FP8 with UE8M0 SF) into the required layout transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights) # Copy inputs into the buffer before each call @@ -146,6 +142,7 @@ The library provides some utility functions besides the above kernels: - `deep_gemm.set_num_sms` / `get_num_sms`: set/get the maximum SM count to use - `deep_gemm.set_tc_util` / `get_tc_util`: set/get an approximated tensor core utilization ratio - `deep_gemm.set_pdl` / `get_pdl`: enable/disable Programmatic Dependent Launch (PDL) +- `deep_gemm.use_deterministic_algorithms`: enable/disable deterministic algorithms - `deep_gemm.set_mk_alignment_for_contiguous_layout` / `get_mk_alignment_for_contiguous_layout`: set/get the group-level M/K alignment for contiguous layout - `deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout`: get the theoretical minimum M/K alignment - `deep_gemm.set_ignore_compile_dims`: configure dimensions to ignore during JIT compilation @@ -158,19 +155,21 @@ The library provides some utility functions besides the above kernels: The library also provides some environment variables, which may be useful: +Each `DG_JIT_*` variable falls back to the corresponding global `DJ_JIT_*` variable when unset. + - General - - `DG_JIT_DEBUG`: `0` or `1`, print JIT debugging information, `0` by default + - `DG_JIT_DEBUG`: `0` or `1`, enable JIT debugging features, including compiler command and PTXAS output, load-time reporting, line info, and PTX/SASS dumps; `0` by default - `DG_PRINT_CONFIGS`: `0` or `1`, print selected configs for each shape, `0` by default - JIT cache - - `DG_JIT_CACHE_DIR`: string, cache directory for compiled kernels, `$HOME/.deep_gemm` by default + - `DG_JIT_CACHE_DIR`: string, cache directory (or a `:`-separated list of directories) for compiled kernels; lookup searches all paths front-to-back (first hit wins) and a cache miss compiles into the first path, `$HOME/.dj` by default - Compiler selection - - `DG_JIT_USE_NVRTC`: `0` or `1`, use NVRTC instead of NVCC (faster compilation, may have lower performance for some cases), `0` by default - - `DG_JIT_NVCC_COMPILER`: string, NVCC compiler path; defaults to `torch.utils.cpp_extension.CUDA_HOME` + - `DG_JIT_NVCC_COMPILER`: string, NVCC compiler path; otherwise CUDA is found through `CUDA_HOME`, `CUDA_PATH`, `which nvcc`, then `/usr/local/cuda` - `DG_JIT_CPP_STANDARD`: integer, C++ standard version, `20` by default - Compiler output - `DG_JIT_PRINT_COMPILER_COMMAND`: `0` or `1`, print compilation commands, `0` by default - `DG_JIT_PTXAS_VERBOSE`: `0` or `1`, show detailed PTXAS output, `0` by default - - `DG_JIT_PTXAS_CHECK`: `0` or `1`, assert no local memory usage in compiled kernels, `0` by default + - `DG_JIT_CHECK_NO_SPILLS`: `0` or `1`, assert no register spills in compiled kernels, `0` by default + - `DG_JIT_CHECK_NO_LOCAL_MEMORY`: `0` or `1`, assert no local memory usage in compiled kernels, `0` by default - `DG_JIT_PRINT_LOAD_TIME`: `0` or `1`, print kernel load time, `0` by default - Debug and profiling - `DG_JIT_WITH_LINEINFO`: `0` or `1`, embed source line info for profiling tools, `0` by default @@ -182,9 +181,8 @@ The library also provides some environment variables, which may be useful: - Build options - `DG_SKIP_CUDA_BUILD`: `0` or `1`, skip CUDA extension build during installation, `0` by default - `DG_FORCE_BUILD`: `0` or `1`, force local build instead of downloading pre-built wheels, `0` by default - - `DG_JIT_USE_RUNTIME_API`: `0` or `1`, use CUDA Runtime API for kernel loading (requires CUDA runtime >= 12.8), `0` by default -For additional examples and details, please refer to [the test code](tests/test_core.py) or review the corresponding Python documentation. +For additional examples and details, please refer to [the test code](tests) or review the corresponding Python documentation. ## Acknowledgement diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 4da25702cc..2f375855ae 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -1,21 +1,23 @@ #pragma once +#include +#include + +#include + #include "../utils/compatibility.hpp" -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm100_mqa_logits.hpp" +#include "../jit_kernels/impls/sm100_sparse_mqa_logits.hpp" #include "../jit_kernels/impls/sm90_fp8_mqa_logits.hpp" -#include "../jit_kernels/impls/smxx_clean_logits.hpp" -#endif #include "layout.hpp" namespace deep_gemm::attention { -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE static void fp8_gemm_nt_skip_head_mid(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -26,7 +28,7 @@ static void fp8_gemm_nt_skip_head_mid(const std::pairget_arch_major(); - const auto epilogue_type = fmt::format("epilogue::transform::EpilogueHeadSplits<{}, {}, {}>", left, mid, right); + const auto arch_major = jit->device.get_arch_major(); + const auto epilogue_type = std::format("epilogue::transform::EpilogueHeadSplits<{}, {}, {}>", left, mid, right); if (arch_major == 9 and sfa.scalar_type() == torch::kFloat and std::get<1>(recipe.value()) != 1) { const auto major_sfb = get_major_type_ab(sfb); sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, std::nullopt, d, m, n, k, major_a, major_b, major_sfb, compiled_dims, epilogue_type); @@ -80,18 +82,18 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple& schedule_meta = std::nullopt) { const auto [q_fp, q_sf] = q; const auto [kv_fp, kv_sf] = kv; const auto qk_dtype = q_fp.scalar_type(); const bool is_fp4 = qk_dtype == kPackedFP4; const bool is_mx_sf = q_sf.has_value(); DG_HOST_ASSERT(not is_fp4 or is_mx_sf); - int seq_len, seq_len_kv, num_heads, head_dim; - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); // Check Q - std::tie(seq_len, num_heads, head_dim) = get_logical_shape<3>(q_fp); + const auto [seq_len, num_heads, head_dim] = get_logical_shape<3>(q_fp); DG_HOST_ASSERT((not is_fp4 and 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() == (is_fp4 ? kPackedFP4 : torch::kFloat8_e4m3fn)); @@ -107,8 +109,7 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple(kv_fp); + const auto [seq_len_kv, _head_dim] = get_logical_shape<2>(kv_fp); DG_HOST_ASSERT(head_dim == _head_dim); DG_HOST_ASSERT(kv_fp.is_contiguous()); DG_HOST_ASSERT(kv_fp.scalar_type() == (is_fp4 ? kPackedFP4 : torch::kFloat8_e4m3fn)); @@ -118,7 +119,7 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple(weights); DG_HOST_ASSERT(seq_len == _seq_len and num_heads == _num_heads); @@ -137,10 +138,9 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple 0 and num_heads <= 128 and num_heads % 4 == 0); constexpr int block_kv = 256; - const int block_q = block_qh / num_heads; - DG_HOST_ASSERT(block_qh % num_heads == 0); + const int block_q = 128 / num_heads; torch::Tensor logits; int aligned_seq_len = align(seq_len, block_q), stride_logits; @@ -158,26 +158,218 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tupledevice.get_arch_major() == 10); + const int num_q_tokens = static_cast(cu_seq_len_k_start.size(0)); + DG_HOST_ASSERT(num_q_tokens > 0 and cu_seq_len_k_end.size(0) == num_q_tokens); + DG_HOST_ASSERT(cu_seq_len_k_start.is_cuda() and cu_seq_len_k_end.is_cuda()); + DG_HOST_ASSERT(cu_seq_len_k_start.is_contiguous() and cu_seq_len_k_end.is_contiguous()); + DG_HOST_ASSERT(cu_seq_len_k_start.scalar_type() == torch::kInt and cu_seq_len_k_end.scalar_type() == torch::kInt); + DG_HOST_ASSERT(num_kv_tokens > 0); + DG_HOST_ASSERT(num_heads > 0 and num_heads <= 128 and num_heads % 4 == 0); + + constexpr int split_kv = 256; + const int block_q = 128 / num_heads; + const int num_sms = runtime->get_num_sms(); + const int required_words = get_mqa_logits_metadata_num_words(num_q_tokens, block_q, num_sms); + const auto schedule_meta = torch::empty({required_words}, cu_seq_len_k_start.options()); + + sm100_mqa_logits_metadata(cu_seq_len_k_start, cu_seq_len_k_end, schedule_meta, + num_q_tokens, num_kv_tokens, block_q, split_kv, num_sms); + return schedule_meta; +} + +static const torch::Tensor& get_sparse_mqa_logits_workspace(const torch::TensorOptions& options, + const int num_q_tokens) { + using namespace layout::sparse_mqa_logits; + constexpr int kNumMaxQTokens = 1 << 20; + constexpr int64_t kNumWorkspaceBytes = + sizeof(WorkspaceState) + static_cast(kNumMaxQTokens) * sizeof(QBlockInfo); + DG_HOST_ASSERT(num_q_tokens <= kNumMaxQTokens); + const auto stream = at::cuda::getCurrentCUDAStream(); + DG_HOST_ASSERT(options.device() == stream.device()); + static std::unordered_map workspaces; + auto& workspace = workspaces[stream]; + if (not workspace.defined()) { + // Warm up each stream before capture so one-time zeroing is not replayed with the graph. + DG_HOST_ASSERT(c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None); + workspace = torch::zeros({kNumWorkspaceBytes}, options.dtype(torch::kByte)); + } + return workspace; +} + +// Each sparse-index row starts with the valid blocks inferred from its KV length. This prefix must +// contain unique, strictly increasing absolute block indices within the corresponding KV range. +// With unaligned ks, block i starts at i * sparse_block_kv + ks % sparse_block_kv. +static torch::Tensor get_sparse_mqa_logits_metadata(const torch::Tensor& cu_seq_len_k_start, + const torch::Tensor& cu_seq_len_k_end, + const int& num_kv_tokens, + const torch::Tensor& sparse_kv_block_indices, + const at::ScalarType& qk_dtype, + const int& sparse_block_kv, + const bool& use_unaligned_ks) { + DG_HOST_ASSERT(jit->device.get_arch_major() == 10); + const auto [num_q_tokens, num_max_sparse_blocks] = get_shape<2>(sparse_kv_block_indices); + DG_HOST_ASSERT(num_q_tokens > 0 and num_kv_tokens >= 0); + DG_HOST_ASSERT(cu_seq_len_k_start.dim() == 1 and cu_seq_len_k_start.size(0) == num_q_tokens); + DG_HOST_ASSERT(cu_seq_len_k_end.dim() == 1 and cu_seq_len_k_end.size(0) == num_q_tokens); + DG_HOST_ASSERT(sparse_kv_block_indices.scalar_type() == torch::kInt and sparse_kv_block_indices.is_contiguous()); + DG_HOST_ASSERT(cu_seq_len_k_start.scalar_type() == torch::kInt and cu_seq_len_k_start.is_contiguous()); + DG_HOST_ASSERT(cu_seq_len_k_end.scalar_type() == torch::kInt and cu_seq_len_k_end.is_contiguous()); + + const int split_kv = get_sparse_mqa_split_kv(qk_dtype); + auto metadata = torch::empty({get_num_metadata_bytes(num_q_tokens, num_max_sparse_blocks, sparse_block_kv, + split_kv, false, runtime->get_num_sms())}, + sparse_kv_block_indices.options().dtype(torch::kUInt8)); + const auto& workspace = get_sparse_mqa_logits_workspace(metadata.options(), num_q_tokens); + launch_sm100_sparse_mqa_logits_metadata(false, use_unaligned_ks, 1, num_kv_tokens, 0, + sparse_kv_block_indices, metadata, workspace, split_kv, sparse_block_kv, + cu_seq_len_k_start.data_ptr(), cu_seq_len_k_end.data_ptr(), + nullptr, nullptr, nullptr); + return metadata; +} + +// Queries belonging to one request must be consecutive. Each sparse-index row starts with unique, +// strictly increasing logical block indices within its context length. Paired queries must also +// have identical block-table rows. +static torch::Tensor get_paged_sparse_mqa_logits_metadata(const torch::Tensor& context_lens, + const torch::Tensor& block_table, + const torch::Tensor& indices, + const int& page_kv, + const torch::Tensor& sparse_kv_block_indices, + const at::ScalarType& qk_dtype, + const int& sparse_block_kv) { + DG_HOST_ASSERT(jit->device.get_arch_major() == 10); + const auto [num_q_tokens, num_max_sparse_blocks] = get_shape<2>(sparse_kv_block_indices); + DG_HOST_ASSERT(num_q_tokens > 0); + DG_HOST_ASSERT(context_lens.numel() == num_q_tokens and context_lens.scalar_type() == torch::kInt and context_lens.is_contiguous()); + DG_HOST_ASSERT(block_table.dim() == 2 and block_table.size(0) == num_q_tokens and block_table.size(1) > 0 and + block_table.scalar_type() == torch::kInt and block_table.stride(1) == 1); + DG_HOST_ASSERT(block_table.stride(0) <= std::numeric_limits::max()); + DG_HOST_ASSERT(indices.dim() == 1 and indices.size(0) == num_q_tokens and + indices.scalar_type() == torch::kInt and indices.is_contiguous()); + DG_HOST_ASSERT(sparse_block_kv == 8 or sparse_block_kv == 16); + DG_HOST_ASSERT(page_kv > 0 and page_kv % sparse_block_kv == 0); + DG_HOST_ASSERT(sparse_kv_block_indices.scalar_type() == torch::kInt and sparse_kv_block_indices.is_contiguous()); + + const int split_kv = get_sparse_mqa_split_kv(qk_dtype); + auto metadata = torch::empty({get_num_metadata_bytes(num_q_tokens, num_max_sparse_blocks, sparse_block_kv, + split_kv, true, runtime->get_num_sms())}, + sparse_kv_block_indices.options().dtype(torch::kUInt8)); + const auto& workspace = get_sparse_mqa_logits_workspace(metadata.options(), num_q_tokens); + launch_sm100_sparse_mqa_logits_metadata(true, false, page_kv, 0, + static_cast(block_table.stride(0)), sparse_kv_block_indices, + metadata, workspace, split_kv, sparse_block_kv, nullptr, nullptr, + context_lens.data_ptr(), block_table.data_ptr(), indices.data_ptr()); + return metadata; +} - // Clean unfilled logits - if (clean_logits) - smxx_clean_logits(logits, cu_seq_len_k_start, cu_seq_len_k_end, 1, seq_len, seq_len_kv, stride_logits); +// Skip metadata header validation to avoid synchronizing the stream +static torch::Tensor fp8_fp4_sparse_mqa_logits(const std::tuple>& q, + const std::tuple& kv, + const torch::Tensor& weights, + const torch::Tensor& metadata, + const int& num_max_sparse_blocks, + const int& sparse_block_kv, + const bool& use_unaligned_ks) { + using namespace layout::sparse_mqa_logits; + const auto [q_fp, q_sf_optional] = q; + const auto [kv_fp, kv_sf] = kv; + DG_HOST_ASSERT(jit->device.get_arch_major() == 10 and q_sf_optional.has_value()); + DG_HOST_ASSERT(num_max_sparse_blocks > 0 and num_max_sparse_blocks % 4 == 0 and num_max_sparse_blocks <= 4096); + DG_HOST_ASSERT(sparse_block_kv == 8 or sparse_block_kv == 16); + const auto& q_sf = q_sf_optional.value(); + const auto qk_dtype = q_fp.scalar_type(); + const bool is_fp4 = qk_dtype == kPackedFP4; + + const auto [num_q_tokens, num_heads, head_dim] = get_logical_shape<3>(q_fp); + DG_HOST_ASSERT(num_q_tokens > 0 and num_heads == kNumHeads and head_dim == kHeadDim); + DG_HOST_ASSERT((is_fp4 or qk_dtype == torch::kFloat8_e4m3fn) and q_fp.is_contiguous()); + const auto [_num_q_tokens_sf, _num_heads_sf] = get_shape<2>(q_sf); + DG_HOST_ASSERT(_num_q_tokens_sf == num_q_tokens and _num_heads_sf == num_heads); + DG_HOST_ASSERT(q_sf.scalar_type() == torch::kInt32 and q_sf.is_contiguous()); + + const auto [num_kv_tokens, kv_head_dim] = get_logical_shape<2>(kv_fp); + DG_HOST_ASSERT(num_kv_tokens > 0 and kv_head_dim == head_dim and kv_fp.scalar_type() == qk_dtype and kv_fp.is_contiguous()); + const auto [_num_kv_tokens_sf] = get_shape<1>(kv_sf); + DG_HOST_ASSERT(_num_kv_tokens_sf == num_kv_tokens and kv_sf.scalar_type() == torch::kInt32 and kv_sf.is_contiguous()); + + const auto [_num_q_tokens_weights, _num_heads_weights] = get_shape<2>(weights); + DG_HOST_ASSERT(_num_q_tokens_weights == num_q_tokens and _num_heads_weights == num_heads); + DG_HOST_ASSERT(weights.scalar_type() == torch::kBFloat16 and weights.stride(1) == 1); + + const int num_output_tokens = num_max_sparse_blocks * sparse_block_kv; + const int logits_stride = align(num_output_tokens, 1024 / static_cast(sizeof(nv_bfloat16))); + auto logits = torch::empty({align(num_q_tokens, kBlockQ), logits_stride}, q_fp.options().dtype(torch::kBFloat16)); + logits = logits.index({torch::indexing::Slice(0, num_q_tokens), torch::indexing::Slice(0, num_output_tokens)}); + launch_sm100_sparse_mqa_logits(false, use_unaligned_ks, sparse_block_kv, + q_fp, q_sf, kv_fp, kv_sf, weights, metadata, logits); + return logits; +} + +// Skip metadata header validation to avoid synchronizing the stream +static torch::Tensor fp8_fp4_paged_sparse_mqa_logits(const std::tuple>& q, + const torch::Tensor& fused_kv_cache, + const torch::Tensor& weights, + const torch::Tensor& metadata, + const int& num_max_sparse_blocks, + const int& sparse_block_kv) { + using namespace layout::sparse_mqa_logits; + const auto [q_fp, q_sf_optional] = q; + DG_HOST_ASSERT(jit->device.get_arch_major() == 10 and q_sf_optional.has_value()); + DG_HOST_ASSERT(num_max_sparse_blocks > 0 and num_max_sparse_blocks % 4 == 0 and num_max_sparse_blocks <= 4096); + DG_HOST_ASSERT(sparse_block_kv == 8 or sparse_block_kv == 16); + const auto& q_sf = q_sf_optional.value(); + const auto qk_dtype = q_fp.scalar_type(); + const bool is_fp4 = qk_dtype == kPackedFP4; + + const auto [num_q_tokens, next_n, num_heads, head_dim] = get_logical_shape<4>(q_fp); + DG_HOST_ASSERT(num_q_tokens > 0 and next_n == 1 and num_heads == kNumHeads and head_dim == kHeadDim); + DG_HOST_ASSERT((is_fp4 or qk_dtype == torch::kFloat8_e4m3fn) and q_fp.is_contiguous()); + const auto [_num_q_tokens_sf, _next_n_sf, _num_heads_sf] = get_shape<3>(q_sf); + DG_HOST_ASSERT(_num_q_tokens_sf == num_q_tokens and _next_n_sf == 1 and _num_heads_sf == num_heads); + DG_HOST_ASSERT(q_sf.scalar_type() == torch::kInt32 and q_sf.is_contiguous()); + + const auto [num_kv_pages, page_kv, num_kv_heads, head_dim_with_sf] = get_shape<4>(fused_kv_cache); + DG_HOST_ASSERT(num_kv_pages > 0 and page_kv > 0 and page_kv % sparse_block_kv == 0 and num_kv_heads == 1 and + head_dim_with_sf == (is_fp4 ? head_dim / 2 : head_dim) + static_cast(sizeof(int))); + DG_HOST_ASSERT(fused_kv_cache.scalar_type() == torch::kUInt8 and fused_kv_cache.stride(1) == head_dim_with_sf and + fused_kv_cache.stride(3) == 1 and fused_kv_cache.stride(0) <= std::numeric_limits::max() and + fused_kv_cache.stride(0) % 512 == 0); + + const auto [_num_q_tokens_weights, _num_heads_weights] = get_shape<2>(weights); + DG_HOST_ASSERT(_num_q_tokens_weights == num_q_tokens and _num_heads_weights == num_heads); + DG_HOST_ASSERT(weights.scalar_type() == torch::kBFloat16 and weights.stride(1) == 1); + + const int num_output_tokens = num_max_sparse_blocks * sparse_block_kv; + const int logits_stride = align(num_output_tokens, 1024 / static_cast(sizeof(nv_bfloat16))); + auto logits = torch::empty({align(num_q_tokens, kBlockQ), logits_stride}, q_fp.options().dtype(torch::kBFloat16)); + logits = logits.index({torch::indexing::Slice(0, num_q_tokens), torch::indexing::Slice(0, num_output_tokens)}); + launch_sm100_sparse_mqa_logits(true, false, sparse_block_kv, q_fp, q_sf, fused_kv_cache, torch::Tensor(), + weights, metadata, logits); return logits; } @@ -195,7 +387,7 @@ static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_ auto schedule_metadata = torch::empty({num_sms + 1, 2}, context_lens.options()); // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (is_varlen) { const auto& indices_tensor = indices.value(); DG_HOST_ASSERT(arch_major == 10 and next_n == 1 and (block_kv == 32 or block_kv == 64 or block_kv == 128)); @@ -231,17 +423,17 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tupleget_num_sms(); - const auto arch_major = device_runtime->get_arch_major(); + int num_sms = runtime->get_num_sms(); + const auto arch_major = jit->device.get_arch_major(); // Check Q - std::tie(batch_size, next_n, num_heads, head_dim) = get_logical_shape<4>(q_fp); + const auto [batch_size, next_n, num_heads, head_dim] = get_logical_shape<4>(q_fp); DG_HOST_ASSERT((not is_fp4 and 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() == (is_fp4 ? kPackedFP4 : torch::kFloat8_e4m3fn)); @@ -258,8 +450,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple(fused_kv_cache); + const auto [num_kv_blocks, block_kv, num_heads_kv, head_dim_with_sf] = get_shape<4>(fused_kv_cache); DG_HOST_ASSERT((arch_major == 10 and (block_kv == 32 or block_kv == 64 or block_kv == 128)) or (arch_major == 9 and block_kv == 64)); const int kv_head_dim = is_fp4 ? head_dim / 2 : head_dim; @@ -283,7 +474,7 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple(weights); DG_HOST_ASSERT(_batch_size_next_n == batch_size * next_n and _num_heads == num_heads); @@ -334,7 +525,6 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tupleset_num_sms(new_num_sms); + runtime->set_num_sms(new_num_sms); }); m.def("get_num_sms", [&]() { - return device_runtime->get_num_sms(); + return runtime->get_num_sms(); }); m.def("set_tc_util", [&](const int& new_tc_util) { - device_runtime->set_tc_util(new_tc_util); + runtime->set_tc_util(new_tc_util); }); m.def("get_tc_util", [&]() { - return device_runtime->get_tc_util(); + return runtime->get_tc_util(); + }); + m.def("set_pdl", [](const bool& new_enable_pdl) { + jit->default_launch_options.enable_pdl = new_enable_pdl; }); - m.def("set_pdl", [&](const bool& new_enable_pdl) { - device_runtime->set_pdl(new_enable_pdl); + m.def("get_pdl", []() { + return *jit->default_launch_options.enable_pdl; }); - m.def("get_pdl", [&]() { - return device_runtime->get_pdl(); + m.def("use_deterministic_algorithms", [&](const bool enabled) { + heuristics_runtime->use_deterministic_algorithms(enabled); }); m.def("set_ignore_compile_dims", [&](const bool& new_value) { heuristics_runtime->set_ignore_compile_dims(new_value); @@ -39,13 +42,6 @@ static void register_apis(pybind11::module_& m) { heuristics_runtime->set_block_size_multiple_of(x, y); } }); - m.def("init", [&](const std::string& library_root_path, const std::string& cuda_home_path_by_python) { -#if DG_TENSORMAP_COMPATIBLE - Compiler::prepare_init(library_root_path, cuda_home_path_by_python); - KernelRuntime::prepare_init(cuda_home_path_by_python); - IncludeParser::prepare_init(library_root_path); -#endif - }); } -} // namespace deep_gemm::runtime +} // namespace deep_gemm::config diff --git a/csrc/apis/einsum.hpp b/csrc/apis/einsum.hpp index becdd2c350..f70e77fae7 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -1,25 +1,24 @@ #pragma once +#include +#include + #include #include #include "../utils/exception.hpp" -#include "../utils/format.hpp" #include "../utils/layout.hpp" #include "../utils/compatibility.hpp" #include "gemm.hpp" -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/sm90_bmk_bnk_mn.hpp" #include "../jit_kernels/impls/sm100_bmk_bnk_mn.hpp" #include "../jit_kernels/impls/sm90_bf16_gemm.hpp" #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" #include "../jit_kernels/impls/smxx_cublaslt.hpp" -#endif namespace deep_gemm::einsum { -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE static void bmk_bnk_mn(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const std::optional& c) { // Currently FP32 only support the accumulated expression @@ -48,7 +47,7 @@ static void bmk_bnk_mn(const torch::Tensor& a, const torch::Tensor& b, const tor DG_HOST_ASSERT(s == s_ and k == k_); // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 9) { sm90_bmn_bnk_mn_gemm(a, b, d, s, m, n, k); } else if (arch_major == 10) { @@ -58,7 +57,7 @@ static void bmk_bnk_mn(const torch::Tensor& a, const torch::Tensor& b, const tor } } -static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D, const bool& use_cublaslt) { +static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D) { const auto [b , h , r ] = get_shape<3>(A); const auto [h_, d , r_] = get_shape<3>(B); const auto [b_, h__, d_] = get_shape<3>(D); @@ -69,8 +68,8 @@ static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const to DG_HOST_ASSERT(D.scalar_type() == torch::kBFloat16 and D.stride(2) == 1); // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); - if (use_cublaslt) { + const auto arch_major = jit->device.get_arch_major(); + if (not heuristics_runtime->get_deterministic_algorithms() and runtime->is_cublaslt_available()) { cublaslt_bhr_hdr_bhd(A, B, D, b, h, r, d); } else if (arch_major == 9) { sm90_bf16_bhr_hdr_bhd(A, B, D, b, h, r, d); @@ -81,7 +80,7 @@ static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const to } } -static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D, const bool& use_cublaslt) { +static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D) { const auto [b , h , d ] = get_shape<3>(A); const auto [h_, d_ , r ] = get_shape<3>(B); const auto [b_, h__, r_] = get_shape<3>(D); @@ -92,8 +91,8 @@ static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const to DG_HOST_ASSERT(D.scalar_type() == torch::kBFloat16 and D.stride(2) == 1); // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); - if (use_cublaslt) { + const auto arch_major = jit->device.get_arch_major(); + if (not heuristics_runtime->get_deterministic_algorithms() and runtime->is_cublaslt_available()) { cublaslt_bhd_hdr_bhr(A, B, D, b, h, r, d); } else if (arch_major == 9) { sm90_bf16_bhd_hdr_bhr(A, B, D, b, h, r, d); @@ -104,12 +103,33 @@ static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const to } } +static void bhd_bhr_hdr(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D, + const std::optional& C) { + const auto [b , h , d ] = get_shape<3>(A); + const auto [b_, h_ , r ] = get_shape<3>(B); + const auto [h__, d_, r_] = get_shape<3>(D); + DG_HOST_ASSERT(b == b_ and h == h_ and h == h__ and d == d_ and r == r_); + + DG_HOST_ASSERT(A.scalar_type() == torch::kBFloat16 and A.stride(2) == 1); + DG_HOST_ASSERT(B.scalar_type() == torch::kBFloat16 and B.stride(2) == 1); + DG_HOST_ASSERT(D.scalar_type() == torch::kFloat and D.stride(2) == 1); + if (C.has_value()) { + DG_HOST_ASSERT(C->scalar_type() == D.scalar_type() and C->sizes() == D.sizes() and C->stride(2) == 1); + } + + // Early return for trivial cases + if (h == 0 or gemm::early_return(d, r, b, D, C)) + return; + + // TODO: investigate cuBLAS determinism. + cublaslt_bhd_bhr_hdr(A, B, D, b, h, r, d, C.has_value()); +} + static void einsum(const std::string& expr, const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, - const std::optional& c, - const bool& use_cublaslt) { + const std::optional& c) { DG_HOST_ASSERT(a.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(b.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat); @@ -121,44 +141,63 @@ static void einsum(const std::string& expr, // TODO: support any expression // TODO: canonicalize expression if (expr == "bmk,bnk->mn") { - DG_HOST_ASSERT(not use_cublaslt); bmk_bnk_mn(a, b, d, c); } else if (expr == "bhr,hdr->bhd") { DG_HOST_ASSERT(not c.has_value()); - bhr_hdr_bhd(a, b, d, use_cublaslt); + bhr_hdr_bhd(a, b, d); } else if (expr == "bhd,hdr->bhr") { DG_HOST_ASSERT(not c.has_value()); - bhd_hdr_bhr(a, b, d, use_cublaslt); + bhd_hdr_bhr(a, b, d); + } else if (expr == "bhd,bhr->hdr") { + bhd_bhr_hdr(a, b, d, c); } else { - DG_HOST_UNREACHABLE(fmt::format("Unsupported einsum expression: {}", expr)); + DG_HOST_UNREACHABLE(std::format("Unsupported einsum expression: {}", expr)); } } +// The D output is either a plain BF16/FP32 tensor, or an FP8 `(d, sfd)` pair +// quantized with dynamic per-32 packed UE8M0 SFs static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, const torch::Tensor& b, const torch::Tensor& sfb, - const torch::Tensor& d, + const std::variant>& d, const std::optional& c, std::optional> recipe, const std::string& compiled_dims) { + const auto d_fp8 = std::get_if>(&d); + const auto d_tensor = d_fp8 != nullptr ? d_fp8->first : std::get(d); + const auto sfd = d_fp8 != nullptr ? std::make_optional(d_fp8->second) : std::nullopt; + // Shape must be `[B, M, K] @ [B, N, K].T` const auto major_a = a.stride(-1) == 1 ? cute::UMMA::Major::K : cute::UMMA::Major::MN; const auto major_b = b.stride(-1) == 1 ? cute::UMMA::Major::K : cute::UMMA::Major::MN; DG_HOST_ASSERT(a.stride(-1) == 1 or a.stride(-2) == 1); DG_HOST_ASSERT(b.stride(-1) == 1 or b.stride(-2) == 1); - DG_HOST_ASSERT(d.stride(-1) == 1); + DG_HOST_ASSERT(d_tensor.stride(-1) == 1); // Type and shape checks const auto [batch_size , m , k ] = get_shape<3>(a); const auto [batch_size_ , n , k_] = get_shape<3>(b); - const auto [batch_size__, m_, n_] = get_shape<3>(d); + const auto [batch_size__, m_, n_] = get_shape<3>(d_tensor); DG_HOST_ASSERT(batch_size == batch_size_ and batch_size == batch_size__); DG_HOST_ASSERT(m == m_ and n == n_ and k == k_); DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat); + if (sfd.has_value()) { + // The SF layout matches a per-token cast of D viewed as a 2D `(m, batch_size * n)` + // matrix, where each row concatenates the `n` columns of all batches. This requires: + // - `d.stride(0) == n`: the batches are contiguous along the columns of the view + // - `n % 128 == 0`: each batch is quantized independently, so neither an SF group + // nor a packed 4-SF word may cross a batch boundary + DG_HOST_ASSERT(jit->device.get_arch_major() == 10 and not c.has_value()); + DG_HOST_ASSERT(d_tensor.scalar_type() == torch::kFloat8_e4m3fn); + DG_HOST_ASSERT(d_tensor.stride(0) == n and n % 128 == 0); + check_sf_layout(sfd.value(), m, batch_size * n, 1, 32, std::nullopt, true, false, torch::kInt); + } else { + DG_HOST_ASSERT(d_tensor.scalar_type() == torch::kBFloat16 or d_tensor.scalar_type() == torch::kFloat); + } // Early return for trivial cases - if (batch_size == 0 or gemm::early_return(m, n, k, d, c)) + if (batch_size == 0 or gemm::early_return(m, n, k, d_tensor, c, sfd)) return; // Transform scaling factors @@ -166,65 +205,71 @@ static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, sfa, sfb, m, n, k, recipe, std::nullopt, std::nullopt, batch_size, batch_size, false); // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 10) { - sm100_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d, batch_size, m, n, k, gran_k_a, gran_k_b, major_a, major_b, compiled_dims); + sm100_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d_tensor, batch_size, m, n, k, gran_k_a, gran_k_b, major_a, major_b, compiled_dims, sfd); } else { const auto major_sfb = get_major_type_ab(sfb); DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == 128); - sm90_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d, batch_size, m, n, k, major_a, major_b, major_sfb, compiled_dims); + sm90_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d_tensor, batch_size, m, n, k, major_a, major_b, major_sfb, compiled_dims); } } static void fp8_einsum(const std::string& expr, const std::pair& a, const std::pair& b, - const torch::Tensor& d, + const std::variant>& d, const std::optional& c, const std::tuple& recipe) { // Some hardcoded Einstein sum kernels - const auto arch_major = device_runtime->get_arch_major(); + // NOTES: only `bhr,hdr->bhd` accepts an FP8 `(d, sfd)` output pair; the other + // expressions take a plain BF16/FP32 D + const auto arch_major = jit->device.get_arch_major(); if (expr == "bhr,hdr->bhd") { // Permute dims to satisfy the order of (batch_size, m, n, k) // (batch_size, m, n, k): (h, b, d, r) + // NOTES: the FP8 output SF columns are flattened over D's last two dims (matching a + // per-token cast of `d.flatten(1, 2)`), so the SF layout is invariant under + // the `(b, h)` permute of the data dims const auto perm_a = a.first.permute({1, 0, 2}); const auto perm_sfa = a.second.permute({1, 0, 2}); - const auto perm_d = d.permute({1, 0, 2}); + auto perm_d = d; + if (const auto d_fp8 = std::get_if>(&perm_d)) + d_fp8->first = d_fp8->first.permute({1, 0, 2}); + else + std::get(perm_d) = std::get(perm_d).permute({1, 0, 2}); const auto perm_c = c.has_value() ? std::make_optional(c.value().permute({1, 0, 2})) : std::nullopt; fp8_bmm(perm_a, perm_sfa, b.first, b.second, perm_d, perm_c, recipe, "nk"); } else if (expr == "bhd,hdr->bhr" and arch_major == 10) { // (batch_size, m, n, k): (h, b, r, d) + DG_HOST_ASSERT(std::holds_alternative(d)); const auto perm_a = a.first.permute({1, 0, 2}); const auto perm_sfa = a.second.permute({1, 0, 2}); const auto perm_b = b.first.permute({0, 2, 1}); const auto perm_sfb = b.second.permute({0, 2, 1}); - const auto perm_d = d.permute({1, 0, 2}); + const auto perm_d = std::get(d).permute({1, 0, 2}); const auto perm_c = c.has_value() ? std::make_optional(c.value().permute({1, 0, 2})) : std::nullopt; fp8_bmm(perm_a, perm_sfa, perm_b, perm_sfb, perm_d, perm_c, recipe, "nk"); } else if (expr == "bhd,bhr->hdr" and arch_major == 10) { // (batch_size, m, n, k): (h, d, r, b) + DG_HOST_ASSERT(std::holds_alternative(d)); const auto perm_a = a.first.permute({1, 2, 0}); const auto perm_sfa = a.second.permute({1, 2, 0}); const auto perm_b = b.first.permute({1, 2, 0}); const auto perm_sfb = b.second.permute({1, 2, 0}); fp8_bmm(perm_a, perm_sfa, perm_b, perm_sfb, d, c, recipe, "mn"); } else { - DG_HOST_UNREACHABLE(fmt::format("Unsupported einsum expression: {}", expr)); + DG_HOST_UNREACHABLE(std::format("Unsupported einsum expression: {}", expr)); } } -#endif - static void register_apis(pybind11::module_& m) { -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE m.def("einsum", &einsum, py::arg("expr"), py::arg("a"), py::arg("b"), - py::arg("d"), py::arg("c") = std::nullopt, - py::arg("use_cublaslt") = false); + py::arg("d"), py::arg("c") = std::nullopt); m.def("fp8_einsum", &fp8_einsum, py::arg("expr"), py::arg("a"), py::arg("b"), - py::arg("d"), py::arg("c") = std::nullopt, + py::arg("d"), py::arg("c") = std::nullopt, py::arg("recipe") = std::make_tuple(1, 128, 128)); -#endif } } // namespace deep_gemm::einsum diff --git a/csrc/apis/gemm.hpp b/csrc/apis/gemm.hpp index 991eabca11..e0b8ec89b9 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -2,13 +2,11 @@ #include "../utils/compatibility.hpp" -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp" #include "../jit_kernels/impls/sm90_bf16_gemm.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" -#endif #include "../jit_kernels/impls/smxx_cublaslt.hpp" @@ -17,16 +15,22 @@ namespace deep_gemm::gemm { static bool early_return(const int& m, const int &n, const int& k, - const torch::Tensor& d, const std::optional& c) { + const torch::Tensor& d, const std::optional& c, + const std::optional& sfd = std::nullopt) { // Do nothing if the problem is empty if (m == 0 or n == 0) return true; // Checks + // NOTES: quantized D never accumulates + DG_HOST_ASSERT(not sfd.has_value() or not c.has_value()); const bool is_cd_same = c.has_value() and c->data_ptr() == d.data_ptr(); if (is_cd_same) DG_HOST_ASSERT(c->sizes() == d.sizes() and c->strides() == d.strides()); - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat); + DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat or + d.scalar_type() == torch::kFloat8_e4m3fn); + // NOTES: an FP8 D must come with SFD, and vice versa; kernels without SFD support must not receive FP8 D + DG_HOST_ASSERT((d.scalar_type() == torch::kFloat8_e4m3fn) == sfd.has_value()); if (c.has_value()) { check_major_type_cd(c.value()); DG_HOST_ASSERT(d.scalar_type() == c.value().scalar_type()); @@ -36,6 +40,9 @@ static bool early_return(const int& m, const int &n, const int& k, if (k == 0) { if (not is_cd_same) c.has_value() ? d.copy_(c.value()) : d.zero_(); + // A zero D dequantizes to zero regardless of the SF content + if (sfd.has_value()) + sfd->zero_(); return true; } @@ -68,8 +75,6 @@ static int check_k_grouped_args(const std::optional>& ks_cpu, return sum_k_if_ks_cpu_missing; } -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE - static void fp8_fp4_gemm_nt(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -78,11 +83,12 @@ static void fp8_fp4_gemm_nt(const std::pair& a, std::optional> recipe_a, std::optional> recipe_b, const std::string& compiled_dims, - const bool& disable_ue8m0_cast) { + const bool& disable_ue8m0_cast, + const std::optional& alpha) { // Shape must be `[M, K] @ [N, K].T` const auto major_a = get_major_type_ab(a.first); const auto major_b = get_major_type_ab(b.first); - if (fp8_requires_k_major()) { + if (fp8_fp4_requires_k_major(a.first, b.first)) { DG_HOST_ASSERT(major_a == cute::UMMA::Major::K); DG_HOST_ASSERT(major_b == cute::UMMA::Major::K); } @@ -91,7 +97,7 @@ static void fp8_fp4_gemm_nt(const std::pair& a, check_major_type_cd(d); // Type and shape checks - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); const auto [m , k ] = check_ab_fp8_fp4(a.first, major_a, arch_major); const auto [n , k_] = check_ab_fp8_fp4(b.first, major_b, arch_major); const auto [m_, n_] = get_shape<2>(d); @@ -108,6 +114,7 @@ static void fp8_fp4_gemm_nt(const std::pair& a, // Dispatch into different implements if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) { + DG_HOST_ASSERT(not alpha.has_value() and "FP8 GEMM alpha requires SM100"); const int gran_n = recipe.has_value() ? std::get<1>(recipe.value()) : std::get<0>(recipe_b.value()); if (gran_n == 1) { sm90_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims); @@ -117,7 +124,7 @@ static void fp8_fp4_gemm_nt(const std::pair& a, } } else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) { sm100_fp8_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, gran_k_a, gran_k_b, - major_a, major_b, compiled_dims); + major_a, major_b, compiled_dims, std::nullopt, alpha); } else { DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types"); } @@ -131,9 +138,10 @@ static void fp8_fp4_gemm_nn(const std::pair& a, const std::optional>& recipe_a, const std::optional>& recipe_b, const std::string& compiled_dims, - const bool& disable_ue8m0_cast) { + const bool& disable_ue8m0_cast, + const std::optional& alpha) { fp8_fp4_gemm_nt(a, {b.first.transpose(0, 1), b.second.transpose(0, 1)}, - d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast); + d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, alpha); } static void fp8_fp4_gemm_tn(const std::pair& a, @@ -144,10 +152,11 @@ static void fp8_fp4_gemm_tn(const std::pair& a, const std::optional>& recipe_a, const std::optional>& recipe_b, const std::string& compiled_dims, - const bool& disable_ue8m0_cast) { + const bool& disable_ue8m0_cast, + const std::optional& alpha) { fp8_fp4_gemm_nt({a.first.transpose(0, 1), a.second.transpose(0, 1)}, {b.first.transpose(0, 1), b.second.transpose(0, 1)}, - d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast); + d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, alpha); } static void fp8_fp4_gemm_tt(const std::pair& a, @@ -158,9 +167,10 @@ static void fp8_fp4_gemm_tt(const std::pair& a, const std::optional>& recipe_a, const std::optional>& recipe_b, const std::string& compiled_dims, - const bool& disable_ue8m0_cast) { + const bool& disable_ue8m0_cast, + const std::optional& alpha) { fp8_fp4_gemm_nt({a.first.transpose(0, 1), a.second.transpose(0, 1)}, b, - d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast); + d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, alpha); } static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair& a, @@ -179,12 +189,12 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pairget_arch_major(); + const auto arch_major = jit->device.get_arch_major(); const auto [m , k ] = check_ab_fp8_fp4(a.first, major_a, arch_major); const auto [num_groups, n, k_] = check_grouped_ab_fp8_fp4(b.first, major_b, arch_major); const auto [m_, n_] = get_shape<2>(d); @@ -264,7 +274,7 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pairget_arch_major(); + const auto arch_major = jit->device.get_arch_major(); const auto [num_groups , m , k ] = check_grouped_ab_fp8_fp4(a.first, major_a, arch_major); const auto [num_groups_ , n , k_] = check_grouped_ab_fp8_fp4(b.first, major_b, arch_major); const auto [num_groups__, m_, n_] = get_shape<3>(d); @@ -308,16 +318,20 @@ static void k_grouped_fp8_gemm_tn_contiguous(const std::pair(recipe) == 1 and std::get<1>(recipe) == 1); + // `k_alignment` must be a multiple of `BLOCK_K = 128` + // All A/B padding must be zero, and the corresponding SF padding must be valid const int gran_k = std::get<2>(recipe); - DG_HOST_ASSERT(gran_k == 32 or gran_k == 128); const int k_alignment = heuristics_runtime->get_mk_alignment_for_contiguous_layout(); - DG_HOST_ASSERT(k_alignment % 32 == 0); + DG_HOST_ASSERT(gran_k == 32 or gran_k == 128); + DG_HOST_ASSERT(k_alignment % 128 == 0); - // Shape checks + // A/B use MN-major layouts `[sum_k, M]` and `[sum_k, N]` const auto [num_groups, m, n] = get_shape<3>(d); const auto [sum_k_ , m_] = get_shape<2>(a.first); const auto [sum_k__, n_] = get_shape<2>(b.first); + // Without PSUM layout, `grouped_layout[i]` is the padded K size of group i; + // With PSUM layout, it is the logical K end, and the group starts at `align(previous_end, k_alignment)` const int sum_k = check_k_grouped_args(ks_cpu, grouped_layout, num_groups, use_psum_layout, k_alignment, static_cast(a.first.size(0))); DG_HOST_ASSERT(m == m_ and n == n_ and sum_k == sum_k_ and sum_k == sum_k__); @@ -325,19 +339,17 @@ static void k_grouped_fp8_gemm_tn_contiguous(const std::pairget_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 10) { + const auto sfa = layout::transform_k_grouped_sf_into_required_layout(a.second, ks_cpu, grouped_layout, recipe, k_alignment, use_psum_layout); + const auto sfb = layout::transform_k_grouped_sf_into_required_layout(b.second, ks_cpu, grouped_layout, recipe, k_alignment, use_psum_layout); sm100_k_grouped_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, grouped_layout, gran_k, k_alignment, cute::UMMA::Major::MN, cute::UMMA::Major::MN, compiled_dims, use_psum_layout); } else { @@ -385,12 +397,12 @@ static void k_grouped_fp8_gemm_nt_contiguous(const std::pairget_num_sms(); + const auto num_sms = runtime->get_num_sms(); const auto tensor_map_buffer = torch::empty({num_sms * 4 * static_cast(sizeof(CUtensorMap))}, a.first.options().dtype(torch::kByte)); // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 9) { sm90_k_grouped_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, ks_cpu.value(), grouped_layout, tensor_map_buffer, cute::UMMA::Major::K, cute::UMMA::Major::K, compiled_dims); @@ -398,14 +410,70 @@ static void k_grouped_fp8_gemm_nt_contiguous(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const std::optional>& ks_cpu, + const torch::Tensor& grouped_layout, + const std::optional& c, + const std::tuple& recipe, + const std::string& compiled_dims, + const bool& use_psum_layout) { + // Must be 1D1D kernel + DG_HOST_ASSERT(recipe == std::make_tuple(1, 1, 32)); + DG_HOST_ASSERT(a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4); + + 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); + + // `k_alignment` must be a multiple of `BLOCK_K = 256` + // All A/B padding codes must be zero, and the corresponding SF padding must be valid + const int k_alignment = heuristics_runtime->get_mk_alignment_for_contiguous_layout(); + DG_HOST_ASSERT(k_alignment % 256 == 0); + + // A/B use K-major logical layouts `[M, sum_k]`/`[N, sum_k]`, + // backed by packed-byte storage `[M, sum_k / 2]`/`[N, sum_k / 2]` + const auto arch_major = jit->device.get_arch_major(); + const auto [num_groups, m, n] = get_shape<3>(d); + const auto [m_, sum_k_] = check_ab_fp8_fp4(a.first, major_a, arch_major); + const auto [n_, sum_k__] = check_ab_fp8_fp4(b.first, major_b, arch_major); + + // Without PSUM layout, `grouped_layout[i]` is the padded K size of group i; + // With PSUM layout, it is the logical K end, and the group starts at `align(previous_end, k_alignment)` + const int sum_k = check_k_grouped_args(ks_cpu, grouped_layout, num_groups, + use_psum_layout, k_alignment, sum_k_); + DG_HOST_ASSERT(m == m_ and n == n_ and sum_k == sum_k_ and sum_k == sum_k__); + + // Contiguity checks + DG_HOST_ASSERT(a.first.is_contiguous()); + DG_HOST_ASSERT(b.first.is_contiguous()); + DG_HOST_ASSERT(d.is_contiguous()); + DG_HOST_ASSERT(not c.has_value() or c.value().is_contiguous()); + + // Early return for trivial cases + if (early_return(m, n, sum_k, d, c)) + return; + + // Transform SF with padding + const auto sfa = layout::transform_k_grouped_sf_into_required_layout(a.second, ks_cpu, grouped_layout, recipe, k_alignment, use_psum_layout); + const auto sfb = layout::transform_k_grouped_sf_into_required_layout(b.second, ks_cpu, grouped_layout, recipe, k_alignment, use_psum_layout); + + // Dispatch implementation + if (arch_major == 10) { + sm100_k_grouped_fp4_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, sum_k, grouped_layout, + k_alignment, major_a, major_b, compiled_dims, use_psum_layout); + } else { + DG_HOST_UNREACHABLE("Unsupported architecture"); + } +} static void bf16_gemm_nt(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const std::optional& c, - const std::string& compiled_dims) { + const std::string& compiled_dims, + const std::optional& alpha) { // Shape must be `[M, K] @ [N, K].T` const auto major_a = get_major_type_ab(a); const auto major_b = get_major_type_ab(b); @@ -427,11 +495,17 @@ static void bf16_gemm_nt(const torch::Tensor& a, return; // Dispatch into different implements - const auto arch_major = device_runtime->get_arch_major(); + if (not heuristics_runtime->get_deterministic_algorithms() and runtime->is_cublaslt_available()) { + cublaslt_gemm(a, b, d, m, n, k, major_a, major_b, c.has_value(), alpha.value_or(1.0f)); + return; + } + + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 9) { + DG_HOST_ASSERT(not alpha.has_value() and "BF16 GEMM alpha requires SM100"); sm90_bf16_gemm(a, b, c, d, m, n, k, major_a, major_b, compiled_dims); } else if (arch_major == 10) { - sm100_bf16_gemm(a, b, c, d, m, n, k, major_a, major_b, compiled_dims); + sm100_bf16_gemm(a, b, c, d, m, n, k, major_a, major_b, compiled_dims, alpha); } else { DG_HOST_UNREACHABLE("Unsupported architecture"); } @@ -441,24 +515,27 @@ static void bf16_gemm_nn(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const std::optional& c, - const std::string& compiled_dims) { - bf16_gemm_nt(a, b.transpose(0, 1), d, c, compiled_dims); + const std::string& compiled_dims, + const std::optional& alpha) { + bf16_gemm_nt(a, b.transpose(0, 1), d, c, compiled_dims, alpha); } static void bf16_gemm_tn(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const std::optional& c, - const std::string& compiled_dims) { - bf16_gemm_nt(a.transpose(0, 1), b.transpose(0, 1), d, c, compiled_dims); + const std::string& compiled_dims, + const std::optional& alpha) { + bf16_gemm_nt(a.transpose(0, 1), b.transpose(0, 1), d, c, compiled_dims, alpha); } static void bf16_gemm_tt(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const std::optional& c, - const std::string& compiled_dims) { - bf16_gemm_nt(a.transpose(0, 1), b, d, c, compiled_dims); + const std::string& compiled_dims, + const std::optional& alpha) { + bf16_gemm_nt(a.transpose(0, 1), b, d, c, compiled_dims, alpha); } static void m_grouped_bf16_gemm_nt_contiguous(const torch::Tensor& a, const torch::Tensor& b, @@ -502,7 +579,7 @@ static void m_grouped_bf16_gemm_nt_contiguous(const torch::Tensor& a, const torc return; // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 9) { sm90_m_grouped_bf16_gemm_contiguous(a, b, d, grouped_layout, num_groups, m, n, k, major_a, major_b, compiled_dims, @@ -551,7 +628,7 @@ static void m_grouped_bf16_gemm_nt_masked(const torch::Tensor& a, const torch::T check_major_type_cd(d); // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 9) { sm90_bf16_m_grouped_gemm_masked(a, b, d, masked_m, num_groups, m, n, k, expected_m, major_a, major_b, compiled_dims); @@ -577,7 +654,7 @@ static void k_grouped_bf16_gemm_tn_contiguous(const torch::Tensor& a, const auto [sum_k__, n_] = get_shape<2>(b); const auto k_alignment = heuristics_runtime->get_mk_alignment_for_contiguous_layout(); - DG_HOST_ASSERT(k_alignment % 32 == 0); + DG_HOST_ASSERT(k_alignment % 128 == 0); const int sum_k = check_k_grouped_args(ks_cpu, grouped_layout, num_groups, use_psum_layout, k_alignment, static_cast(a.size(0))); DG_HOST_ASSERT(m == m_ and n == n_ and sum_k == sum_k_ and sum_k == sum_k__); @@ -586,16 +663,17 @@ static void k_grouped_bf16_gemm_tn_contiguous(const torch::Tensor& a, DG_HOST_ASSERT(a.is_contiguous()); DG_HOST_ASSERT(b.is_contiguous()); DG_HOST_ASSERT(d.is_contiguous()); - DG_HOST_ASSERT(c.has_value() and c.value().is_contiguous()); + DG_HOST_ASSERT(not c.has_value() or c.value().is_contiguous()); // Early return for trivial cases if (early_return(m, n, sum_k, d, c)) return; // Dispatch implementation - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 9) { - // No psum on SM90 + // No direct output or psum on SM90 + DG_HOST_ASSERT(c.has_value()); DG_HOST_ASSERT(not use_psum_layout and ks_cpu.has_value() and not ks_cpu.value().empty()); sm90_bf16_k_grouped_gemm(a, b, c, d, m, n, ks_cpu.value(), grouped_layout, cute::UMMA::Major::MN, cute::UMMA::Major::MN, compiled_dims); @@ -606,8 +684,6 @@ static void k_grouped_bf16_gemm_tn_contiguous(const torch::Tensor& a, DG_HOST_UNREACHABLE("Unsupported architecture"); } } -#endif - static void cublaslt_gemm_nt(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, const std::optional& c) { // Shape must be `[M, K] @ [N, K].T` @@ -642,34 +718,115 @@ static void cublaslt_gemm_tt(const torch::Tensor& a, const torch::Tensor& b, cublaslt_gemm_nt(a.transpose(0, 1), b, d, c); } +static void cublaslt_nvfp4_gemm_nt(const std::pair& a, + const std::pair& b, + const torch::Tensor& d, + const std::optional& c) { + // Shape must be `[M, K] @ [N, K].T` with both operands packed FP4 and K-major + DG_HOST_ASSERT(a.first.scalar_type() == kPackedFP4 and b.first.scalar_type() == kPackedFP4); + DG_HOST_ASSERT(a.first.is_contiguous() and b.first.is_contiguous()); + + // Type and shape checks + // NOTES: shapes are in logical FP4 elements (2 elements per packed byte) + const auto [m , packed_k ] = get_shape<2>(a.first); + const auto [n , packed_k_] = get_shape<2>(b.first); + const auto [m_, n_] = get_shape<2>(d); + const auto k = packed_k * 2; + DG_HOST_ASSERT(m == m_ and n == n_ and packed_k == packed_k_); + DG_HOST_ASSERT(k % 32 == 0); + + // Scaling factors must be UE4M3 bytes (NVFP4 recipe, 16-element blocks) in the + // cuBLASLt tiled layout: `[ceil(mn / 128), ceil(k / 64), 32, 4, 4]` in bytes + const auto& check_sf = [&](const torch::Tensor& sf, const int& mn) { + DG_HOST_ASSERT(sf.scalar_type() == torch::kByte or sf.scalar_type() == torch::kChar); + DG_HOST_ASSERT(sf.is_contiguous()); + DG_HOST_ASSERT(sf.numel() == static_cast(ceil_div(mn, 128)) * ceil_div(k, 64) * 512); + }; + check_sf(a.second, m); + check_sf(b.second, n); + + // Early return for trivial cases + if (early_return(m, n, k, d, c)) + return; + + cublaslt_nvfp4_gemm(a.first, a.second, b.first, b.second, d, m, n, k, c.has_value()); +} + +static auto get_cublaslt_batched_view(const torch::Tensor& tensor) { + DG_HOST_ASSERT(tensor.dim() == 2 or tensor.dim() == 3); + return tensor.dim() == 2 ? tensor.unsqueeze(0) : tensor; +} + +static void batched_syrk(const torch::Tensor& a, const torch::Tensor& d) { + // D = A @ A.mT. A and D may be either unbatched 2D tensors or batched 3D tensors. + const auto a_3d = get_cublaslt_batched_view(a); + const auto d_3d = get_cublaslt_batched_view(d); + const auto [num_batches, m, k] = get_shape<3>(a_3d); + const auto [num_batches_, m_, n_] = get_shape<3>(d_3d); + const auto major_a = get_major_type_ab(a_3d); + check_major_type_cd(d_3d); + DG_HOST_ASSERT(num_batches == num_batches_ and m == m_ and m == n_); + DG_HOST_ASSERT(a_3d.scalar_type() == d_3d.scalar_type()); + + if (num_batches == 0 or early_return(m, m, k, d_3d, std::nullopt)) + return; + cublaslt_batched_gemm(a_3d, a_3d, d_3d, m, m, k, num_batches, major_a, major_a); +} + +static void batched_symm(const torch::Tensor& a, const torch::Tensor& b, + const torch::Tensor& d) { + // D = A @ B. A is symmetric by contract; cuBLASLt executes this as strided-batched GEMM. + const auto a_3d = get_cublaslt_batched_view(a); + const auto b_3d = get_cublaslt_batched_view(b); + const auto d_3d = get_cublaslt_batched_view(d); + const auto [num_batches, m, m_] = get_shape<3>(a_3d); + const auto [num_batches_b, m_b, k] = get_shape<3>(b_3d); + const auto [num_batches_d, m_d, k_d] = get_shape<3>(d_3d); + const auto major_a = get_major_type_ab(a_3d); + check_major_type_cd(d_3d); + DG_HOST_ASSERT(num_batches == num_batches_b and num_batches == num_batches_d); + DG_HOST_ASSERT(m == m_ and m == m_b and m == m_d and k == k_d); + DG_HOST_ASSERT(a_3d.scalar_type() == b_3d.scalar_type() and a_3d.scalar_type() == d_3d.scalar_type()); + + if (num_batches == 0 or early_return(m, k, m, d_3d, std::nullopt)) + return; + + const auto b_transposed = b_3d.transpose(1, 2); + const auto major_b = get_major_type_ab(b_transposed); + cublaslt_batched_gemm(a_3d, b_transposed, d_3d, m, k, m, num_batches, major_a, major_b); +} + static void register_apis(pybind11::module_& m) { -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE // FP8 FP4 GEMMs m.def("fp8_fp4_gemm_nt", &fp8_fp4_gemm_nt, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt, 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("disable_ue8m0_cast") = false, + py::arg("alpha") = std::nullopt); m.def("fp8_fp4_gemm_nn", &fp8_fp4_gemm_nn, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt, 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("disable_ue8m0_cast") = false, + py::arg("alpha") = std::nullopt); m.def("fp8_fp4_gemm_tn", &fp8_fp4_gemm_tn, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt, py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, py::arg("compiled_dims") = "mn", - py::arg("disable_ue8m0_cast") = false); + py::arg("disable_ue8m0_cast") = false, + py::arg("alpha") = std::nullopt); m.def("fp8_fp4_gemm_tt", &fp8_fp4_gemm_tt, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt, py::arg("recipe_a") = std::nullopt, py::arg("recipe_b") = std::nullopt, py::arg("compiled_dims") = "mn", - py::arg("disable_ue8m0_cast") = false); + py::arg("disable_ue8m0_cast") = false, + py::arg("alpha") = std::nullopt); m.def("m_grouped_fp8_fp4_gemm_nt_contiguous", &m_grouped_fp8_fp4_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("recipe") = std::nullopt, @@ -706,35 +863,46 @@ static void register_apis(pybind11::module_& m) { py::arg("recipe") = std::make_tuple(1, 1, 128), py::arg("compiled_dims") = "mn", py::arg("use_psum_layout") = false); + m.def("k_grouped_fp4_gemm_nt_contiguous", &k_grouped_fp4_gemm_nt_contiguous, + py::arg("a"), py::arg("b"), py::arg("d"), + py::arg("ks_cpu"), py::arg("grouped_layout"), + py::arg("c") = std::nullopt, + py::arg("recipe") = std::make_tuple(1, 1, 32), + py::arg("compiled_dims") = "mn", + py::arg("use_psum_layout") = false); - // FP8 GEMM alias names + // FP4 and FP8 GEMM alias names + m.attr("fp4_gemm_nt") = m.attr("fp8_fp4_gemm_nt"); m.attr("fp8_gemm_nt") = m.attr("fp8_fp4_gemm_nt"); m.attr("fp8_gemm_nn") = m.attr("fp8_fp4_gemm_nn"); m.attr("fp8_gemm_tn") = m.attr("fp8_fp4_gemm_tn"); m.attr("fp8_gemm_tt") = m.attr("fp8_fp4_gemm_tt"); + m.attr("m_grouped_fp4_gemm_nt_contiguous") = m.attr("m_grouped_fp8_fp4_gemm_nt_contiguous"); m.attr("m_grouped_fp8_gemm_nt_contiguous") = m.attr("m_grouped_fp8_fp4_gemm_nt_contiguous"); m.attr("m_grouped_fp8_gemm_nn_contiguous") = m.attr("m_grouped_fp8_fp4_gemm_nn_contiguous"); + m.attr("m_grouped_fp4_gemm_nt_masked") = m.attr("m_grouped_fp8_fp4_gemm_nt_masked"); m.attr("m_grouped_fp8_gemm_nt_masked") = m.attr("m_grouped_fp8_fp4_gemm_nt_masked"); -#endif - -#if DG_TENSORMAP_COMPATIBLE // BF16 GEMMs m.def("bf16_gemm_nt", &bf16_gemm_nt, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "nk"); + py::arg("compiled_dims") = "nk", + py::arg("alpha") = std::nullopt); m.def("bf16_gemm_nn", &bf16_gemm_nn, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "nk"); + py::arg("compiled_dims") = "nk", + py::arg("alpha") = std::nullopt); m.def("bf16_gemm_tn", &bf16_gemm_tn, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "mn"); + py::arg("compiled_dims") = "mn", + py::arg("alpha") = std::nullopt); m.def("bf16_gemm_tt", &bf16_gemm_tt, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt, - py::arg("compiled_dims") = "mn"); + py::arg("compiled_dims") = "mn", + py::arg("alpha") = std::nullopt); m.def("m_grouped_bf16_gemm_nt_contiguous", &m_grouped_bf16_gemm_nt_contiguous, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("grouped_layout"), py::arg("compiled_dims") = "nk", @@ -755,8 +923,6 @@ static void register_apis(pybind11::module_& m) { py::arg("c") = std::nullopt, py::arg("compiled_dims") = "mn", py::arg("use_psum_layout") = false); -#endif - // cuBLASLt GEMMs m.def("cublaslt_gemm_nt", &cublaslt_gemm_nt, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt); @@ -766,6 +932,12 @@ static void register_apis(pybind11::module_& m) { py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt); m.def("cublaslt_gemm_tt", &cublaslt_gemm_tt, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt); + m.def("cublaslt_nvfp4_gemm_nt", &cublaslt_nvfp4_gemm_nt, + py::arg("a"), py::arg("b"), py::arg("d"), py::arg("c") = std::nullopt); + m.def("batched_syrk", &batched_syrk, + py::arg("a"), py::arg("d")); + m.def("batched_symm", &batched_symm, + py::arg("a"), py::arg("b"), py::arg("d")); } } // namespace deep_gemm::gemm diff --git a/csrc/apis/hyperconnection.hpp b/csrc/apis/hyperconnection.hpp index 1a13984d24..c23bbc9a64 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -2,14 +2,11 @@ #include "../utils/compatibility.hpp" -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp" #include "../jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp" -#endif namespace deep_gemm::hyperconnection { -#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE static void tf32_hc_prenorm_gemm(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, @@ -47,7 +44,7 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a, return; // Dispatch into different implements - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 9) { sm90_tf32_hc_prenorm_gemm(a, b, d, sqr_sum, m, n, k, num_splits.has_value() ? num_splits.value() : 1); } else if (arch_major == 10) { @@ -57,14 +54,11 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a, } } -#endif 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, py::arg("a"), py::arg("b"), py::arg("d"), py::arg("sqr_sum"), py::arg("num_splits") = std::nullopt); -#endif } } // namespace deep_gemm::hyperconnection diff --git a/csrc/apis/layout.hpp b/csrc/apis/layout.hpp index 800484de0e..814ebdec5a 100644 --- a/csrc/apis/layout.hpp +++ b/csrc/apis/layout.hpp @@ -4,13 +4,10 @@ #include "../utils/layout.hpp" #include "../utils/compatibility.hpp" -#if DG_TENSORMAP_COMPATIBLE #include "../jit_kernels/impls/smxx_layout.hpp" -#endif namespace deep_gemm::layout { -#if DG_TENSORMAP_COMPATIBLE static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, const int& mn, const int& k, const std::variant, @@ -19,7 +16,7 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf, const std::optional& is_sfa, const bool& disable_ue8m0_cast, const std::optional& psum_layout = std::nullopt) { - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); // Get granularity MN/K from recipe int gran_mn, gran_k; @@ -99,10 +96,9 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te DG_HOST_ASSERT(std::get<0>(recipe) == 1 and std::get<1>(recipe) == 1); const int gran_k = std::get<2>(recipe); - DG_HOST_ASSERT(gran_k == 32 or gran_k == 128); - DG_HOST_ASSERT(k_alignment % 32 == 0); - - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); + DG_HOST_ASSERT((arch_major == 9 and gran_k == 128 and k_alignment == 128) or + (arch_major == 10 and (gran_k == 32 or gran_k == 128) and k_alignment % 128 == 0)); // FP32 on SM90 if (sf.scalar_type() == torch::kFloat and arch_major == 9) @@ -112,17 +108,14 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te if (sf.scalar_type() == torch::kFloat and arch_major == 10) return get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(sf, grouped_layout, ks_cpu, gran_k, k_alignment, use_psum_layout); - // INT (already packed UE8M0) on SM100 - if (sf.scalar_type() == torch::kInt and arch_major == 10) + // Pre-packed UE8M0 is only accepted for gran_k=32 on SM100. + if (sf.scalar_type() == torch::kInt and arch_major == 10 and gran_k == 32) return check_k_grouped_packed_ue8m0_tensor(sf, grouped_layout, ks_cpu, gran_k, k_alignment, use_psum_layout); DG_HOST_UNREACHABLE("Unknown cases"); } -#endif - static void register_apis(pybind11::module_& m) { -#if DG_TENSORMAP_COMPATIBLE m.def("transform_sf_into_required_layout", &transform_sf_into_required_layout, py::arg("sf"), py::arg("mn"), py::arg("k"), py::arg("recipe"), py::arg("num_groups") = std::nullopt, @@ -137,8 +130,6 @@ static void register_apis(pybind11::module_& m) { m.def("get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", &get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor, py::arg("sf"), py::arg("grouped_layout"), py::arg("ks_cpu"), py::arg("gran_k"), py::arg("k_alignment"), py::arg("use_psum_layout") = false); -#endif - m.def("set_mk_alignment_for_contiguous_layout", [&](const int& new_value) { heuristics_runtime->set_mk_alignment_for_contiguous_layout(new_value); }); diff --git a/csrc/apis/mega_gate.hpp b/csrc/apis/mega_gate.hpp new file mode 100644 index 0000000000..617a43fb81 --- /dev/null +++ b/csrc/apis/mega_gate.hpp @@ -0,0 +1,234 @@ +#pragma once + +#include +#include +#include +#include +#include + +#include +#include + +#include "../runtime/runtime.hpp" +#include "../jit_kernels/impls/sm100_bf16_mega_gate.hpp" +#include "../utils/exception.hpp" +#include "../utils/layout.hpp" + +namespace deep_gemm::mega_gate { + +static int get_scoring_type(const std::string& scoring_func) { + if (scoring_func == "sigmoid") + return static_cast(mega_gate_layout::ScoringType::Sigmoid); + if (scoring_func == "sqrtsoftplus") + return static_cast(mega_gate_layout::ScoringType::SqrtSoftplus); + if (scoring_func == "identity") + return static_cast(mega_gate_layout::ScoringType::Identity); + DG_HOST_UNREACHABLE("Unsupported MoE scoring function"); +} + +static const torch::Tensor& get_score_barriers(const torch::TensorOptions& options) { + const auto stream = at::cuda::getCurrentCUDAStream(); + DG_HOST_ASSERT(options.device() == stream.device()); + static std::unordered_map score_barriers_by_stream; + auto& score_barriers = score_barriers_by_stream[stream]; + if (not score_barriers.defined()) { + DG_HOST_ASSERT(c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None); + score_barriers = torch::zeros({mega_gate_layout::kNumMaxTokenBlocks, mega_gate_layout::kScoreBarrierLineBytes}, + options.dtype(torch::kByte)); + } + return score_barriers; +} + +static void check_same_cuda_device(const torch::Tensor& tensor, const torch::Tensor& reference) { + DG_HOST_ASSERT(tensor.is_cuda() and tensor.get_device() == reference.get_device()); +} + +static void check_tensor(const torch::Tensor& tensor, const torch::Tensor& reference, + const at::IntArrayRef& shape, const torch::ScalarType& scalar_type, + const bool& contiguous = true) { + check_same_cuda_device(tensor, reference); + DG_HOST_ASSERT(tensor.sizes() == shape and tensor.scalar_type() == scalar_type and + (not contiguous or tensor.is_contiguous())); +} + +/** + * Fuse the BF16 gate GEMM with the MoE top-k gate. + * + * Args: + * x [T,H] BF16 + * weight [E,H] BF16 + * num_topk int + * use_shared_as_routed bool + * num_shared_experts int + * routed_scaling_factor float + * ep_rank int + * scoring_func "sigmoid" | "sqrtsoftplus" | "identity" + * mask [T] bool | None + * bias [E] FP32 | None + * image_bias [E] FP32 | None + * image_token_mask [T] bool | None + * fix_routing_mask [T] bool | None + * to_physical_map [E+S, dup] int32 | None + * logical_count [E+S] int32 | None + * unmapped_topk_idx [T,K] int64 | None + * force_random [T] bool | None + * out ([T,K'] int64, [T,K'] FP32) | None + * + * Returns: + * (topk_idx, topk_weights) + * + * Notes: + * S = num_shared_experts when use_shared_as_routed, else 0; K' = K + S. + * The kernel ranks experts on score + bias, while the emitted weights are the + * unbiased scores normalized over the top-k sum and scaled by routed_scaling_factor. + */ +static std::tuple +bf16_mega_gate(const torch::Tensor& x, + const torch::Tensor& weight, + const int& num_topk, + const bool& use_shared_as_routed, + const int& num_shared_experts, + const float& routed_scaling_factor, + const int& ep_rank, + const std::string& scoring_func, + const std::optional& mask, + const std::optional& bias, + const std::optional& image_bias, + const std::optional& image_token_mask, + const std::optional& fix_routing_mask, + const std::optional& to_physical_map, + const std::optional& logical_count, + const std::optional& unmapped_topk_idx, + const std::optional& force_random, + const std::optional>& out) { + const auto [num_tokens, hidden] = get_shape<2>(x); + const auto [num_routed_experts, weight_hidden] = get_shape<2>(weight); + DG_HOST_ASSERT(hidden == weight_hidden); + DG_HOST_ASSERT(hidden > 0 and hidden % 256 == 0); + DG_HOST_ASSERT(num_routed_experts > 0 and num_routed_experts <= 512 and num_routed_experts % 4 == 0); + DG_HOST_ASSERT(x.scalar_type() == torch::kBFloat16 and weight.scalar_type() == torch::kBFloat16 and + x.is_contiguous() and weight.is_contiguous()); + check_same_cuda_device(weight, x); + + DG_HOST_ASSERT(num_topk > 0 and num_topk <= num_routed_experts and num_topk <= 32); + DG_HOST_ASSERT(ep_rank >= 0); + DG_HOST_ASSERT(std::isfinite(routed_scaling_factor)); + const auto scoring_type = get_scoring_type(scoring_func); + + int effective_num_shared_experts = 0; + if (use_shared_as_routed) { + DG_HOST_ASSERT(num_shared_experts == 1 or num_shared_experts == 2); + DG_HOST_ASSERT(num_topk % num_shared_experts == 0); + DG_HOST_ASSERT(num_routed_experts % (num_topk / num_shared_experts) == 0); + effective_num_shared_experts = num_shared_experts; + } + const auto num_physical_topk = num_topk + effective_num_shared_experts; + DG_HOST_ASSERT(num_physical_topk <= 32); + + if (mask.has_value()) + check_tensor(mask.value(), x, {num_tokens}, torch::kBool); + if (bias.has_value()) + check_tensor(bias.value(), x, {num_routed_experts}, torch::kFloat32); + + DG_HOST_ASSERT(image_bias.has_value() == image_token_mask.has_value()); + if (image_bias.has_value()) { + check_tensor(image_bias.value(), x, {num_routed_experts}, torch::kFloat32); + check_tensor(image_token_mask.value(), x, {num_tokens}, torch::kBool); + } + + DG_HOST_ASSERT(to_physical_map.has_value() == logical_count.has_value()); + const auto num_logical_experts = num_routed_experts + effective_num_shared_experts; + if (to_physical_map.has_value()) { + const auto& physical_map = to_physical_map.value(); + DG_HOST_ASSERT(physical_map.dim() == 2); + DG_HOST_ASSERT(physical_map.size(1) > 0); + check_tensor(physical_map, x, {num_logical_experts, physical_map.size(1)}, torch::kInt32); + check_tensor(logical_count.value(), x, {num_logical_experts}, torch::kInt32); + } + + if (unmapped_topk_idx.has_value()) { + const auto& unmapped = unmapped_topk_idx.value(); + check_tensor(unmapped, x, {num_tokens, num_topk}, torch::kInt64, false); + DG_HOST_ASSERT(unmapped.stride(1) == 1); + } + if (fix_routing_mask.has_value()) { + DG_HOST_ASSERT(unmapped_topk_idx.has_value()); + check_tensor(fix_routing_mask.value(), x, {num_tokens}, torch::kBool); + } + if (force_random.has_value()) + check_tensor(force_random.value(), x, {num_tokens}, torch::kBool); + + torch::Tensor topk_idx, topk_weights; + if (out.has_value()) { + std::tie(topk_idx, topk_weights) = out.value(); + check_tensor(topk_idx, x, {num_tokens, num_physical_topk}, torch::kInt64); + check_tensor(topk_weights, x, {num_tokens, num_physical_topk}, torch::kFloat32); + } else { + topk_idx = torch::empty({num_tokens, num_physical_topk}, x.options().dtype(torch::kInt64)); + topk_weights = torch::empty({num_tokens, num_physical_topk}, x.options().dtype(torch::kFloat32)); + } + + if (num_tokens == 0) + return {topk_idx, topk_weights}; + DG_HOST_ASSERT(num_tokens <= static_cast(layout::mega_gate::kNumMaxTokens)); + + const auto arch_major = jit->device.get_arch_major(); + DG_HOST_ASSERT(arch_major == 10); + const auto config = get_sm100_bf16_mega_gate_config(num_tokens, hidden, num_routed_experts, + runtime->get_num_sms(), bias.has_value(), + image_token_mask.has_value(), to_physical_map.has_value()); + const auto num_aligned_experts = align(num_routed_experts, static_cast(mega_gate_layout::kExpertAlignment)); + const auto num_token_blocks = ceil_div(num_tokens, config.block_tokens); + const auto num_scratch_bytes = mega_gate_layout::Workspace<>::get_num_scratch_bytes(num_token_blocks, config.num_split_k, + config.block_tokens, num_aligned_experts); + const auto scratch = torch::empty({static_cast(num_scratch_bytes)}, x.options().dtype(torch::kByte)); + const auto& score_barriers = get_score_barriers(x.options()); + sm100_bf16_mega_gate(x, weight, bias, image_bias, image_token_mask, mask, fix_routing_mask, + to_physical_map, logical_count, topk_idx, unmapped_topk_idx, + topk_weights, force_random, num_tokens, hidden, num_routed_experts, + num_topk, effective_num_shared_experts, routed_scaling_factor, + ep_rank, scoring_type, config, scratch, score_barriers); + return {topk_idx, topk_weights}; +} + +static pybind11::dict get_bf16_mega_gate_config(const int& num_tokens, const int& hidden, + const int& num_routed_experts, const int& num_topk) { + DG_HOST_ASSERT(num_tokens > 0 and num_tokens <= static_cast(layout::mega_gate::kNumMaxTokens)); + DG_HOST_ASSERT(hidden > 0 and hidden % 256 == 0); + DG_HOST_ASSERT(num_routed_experts > 0 and num_routed_experts <= 512 and num_routed_experts % 4 == 0); + DG_HOST_ASSERT(num_topk > 0 and num_topk <= num_routed_experts and num_topk <= 32); + const auto num_device_sms = runtime->get_num_sms(); + const auto config = get_sm100_bf16_mega_gate_config(num_tokens, hidden, num_routed_experts, + num_device_sms, true, true, true); + pybind11::dict result; + result["block_tokens"] = config.block_tokens; + result["num_mma_ctas"] = config.num_mma_ctas; + result["num_split_k"] = config.num_split_k; + result["num_expert_groups"] = config.num_expert_groups; + result["num_gate_warpgroups"] = config.num_gate_warpgroups; + result["num_sms"] = config.num_launch_sms; + return result; +} + +static void register_apis(pybind11::module_& m) { + m.def("get_bf16_mega_gate_config", &get_bf16_mega_gate_config, + py::arg("num_tokens"), py::arg("hidden"), + py::arg("num_routed_experts"), py::arg("num_topk")); + m.def("bf16_mega_gate", &bf16_mega_gate, + py::arg("x"), py::arg("weight"), py::arg("num_topk"), + py::arg("use_shared_as_routed"), py::arg("num_shared_experts"), + py::arg("routed_scaling_factor"), py::arg("ep_rank"), + py::arg("scoring_func") = "identity", + py::arg("mask") = std::nullopt, + py::arg("bias") = std::nullopt, + py::arg("image_bias") = std::nullopt, + py::arg("image_token_mask") = std::nullopt, + py::arg("fix_routing_mask") = std::nullopt, + py::arg("to_physical_map") = std::nullopt, + py::arg("logical_count") = std::nullopt, + py::arg("unmapped_topk_idx") = std::nullopt, + py::arg("force_random") = std::nullopt, + py::arg("out") = std::nullopt); +} + +} // namespace deep_gemm::mega_gate diff --git a/csrc/apis/mega_mhc.hpp b/csrc/apis/mega_mhc.hpp new file mode 100644 index 0000000000..1cbc0c8230 --- /dev/null +++ b/csrc/apis/mega_mhc.hpp @@ -0,0 +1,287 @@ +#pragma once + +#include +#include + +#include "../utils/compatibility.hpp" + +#include "../jit_kernels/impls/sm100_mega_mhc.hpp" +#include "../utils/layout.hpp" + +namespace deep_gemm::mega_mhc { + +namespace mhc_layout = layout::mega_mhc; + +static uint32_t get_num_splits(const uint32_t num_m_blocks, const uint32_t hidden, const uint32_t num_sms) { + if (heuristics_runtime->get_deterministic_algorithms()) + return mhc_layout::kDefaultNumSplits; + + // Fewest splits that still reach the shortest possible longest task + const uint32_t num_k_blocks = hidden / mhc_layout::BLOCK_K; + const uint32_t max_num_splits = cute::min( + mhc_layout::kNumMaxSplits, cute::max(mhc_layout::kDefaultNumSplits, num_sms / num_m_blocks)); + return math::ceil_div(num_k_blocks, math::ceil_div(num_k_blocks, max_num_splits)); +} + +static const torch::Tensor& get_split_barriers(const torch::TensorOptions& options) { + const auto stream = at::cuda::getCurrentCUDAStream(); + DG_HOST_ASSERT(options.device() == stream.device()); + static std::unordered_map split_barriers_by_stream; + auto& split_barriers = split_barriers_by_stream[stream]; + if (not split_barriers.defined()) { + // Warm up each stream before capture so one-time zeroing is not replayed with the graph. + DG_HOST_ASSERT(c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None); + split_barriers = torch::zeros( + {mhc_layout::kNumSplitBarriers, mhc_layout::kNumMaxMBlocks, mhc_layout::kSplitBarrierLineBytes}, + options.dtype(torch::kByte)); + } + return split_barriers; +} + +/** + * Execute shifted or normal mHC, generate the next mixes, and apply RMSNorm + * with optional FP8 E4M3 casting. + * + * Args: + * x [T,H] BF16 + * residual [T,hc_mult,H] BF16 + * shifted_prev_mix [T,hc_mult,1] FP32 | None + * post_mix [T,hc_mult,1] FP32 + * comb_res_mix [T,hc_mult,hc_mult] FP32 + * fn [hc_mult*(hc_mult+2),hc_mult*H] FP32 + * mix_scales [3] FP32 + * mix_bases [hc_mult*(hc_mult+2)] FP32 + * hc_mult int + * hc_norm_eps float + * hc_pre_eps float + * hc_post_scale float + * sinkhorn_eps float + * num_sinkhorn_iters int + * rmsnorm_weight [H] BF16 + * rmsnorm_eps float + * rmsnorm_scale float + * new_residual [T,hc_mult,H] BF16 + * new_prev_mix [T,hc_mult,1] FP32 | None + * new_post_mix [T,hc_mult,1] FP32 + * new_comb_res_mix [T,hc_mult,hc_mult] FP32 + * y_bf16 [T,H] BF16 | None + * y_fp8 [T,H] E4M3 | None + * y_gemm_sf [T,H/128] int32 | None + * y_routed_sf [T,H/128] int32 | None + * y_shared_sf [T,H/128] int32 | None + * shared_sf_block_m int + * + * Returns: + * None + * + * Notes: + * shifted_prev_mix and new_prev_mix are both tensors for shifted mHC and + * both None for normal mHC. Split-K is selected per invocation: 16 with + * deterministic algorithms; otherwise the fewest unequal whole-block splits + * that reach the shortest longest producer task. FP8 gran_k is 32, with four + * UE8M0 scaling factors packed into each int32. + * y_bf16 and y_fp8 are independent optional outputs; at least one is required. + * With y_fp8, supply exactly either y_gemm_sf alone, or y_routed_sf and + * y_shared_sf together. y_gemm_sf is TMA-aligned column-major, + * y_routed_sf is contiguous row-major, and y_shared_sf is the Mega MoE + * shared-expert layout and requires shared_sf_block_m > 0. + */ +static void mega_mhc(const torch::Tensor& x, + const torch::Tensor& residual, + const std::optional& shifted_prev_mix, + const torch::Tensor& post_mix, + const torch::Tensor& comb_res_mix, + const torch::Tensor& fn, + const torch::Tensor& mix_scales, + const torch::Tensor& mix_bases, + const int& hc_mult, + const float& hc_norm_eps, + const float& hc_pre_eps, + const float& hc_post_scale, + const float& sinkhorn_eps, + const int& num_sinkhorn_iters, + const torch::Tensor& rmsnorm_weight, + const float& rmsnorm_eps, + const float& rmsnorm_scale, + const torch::Tensor& new_residual, + const std::optional& new_prev_mix, + const torch::Tensor& new_post_mix, + const torch::Tensor& new_comb_res_mix, + const std::optional& y_bf16, + const std::optional& y_fp8, + const std::optional& y_gemm_sf, + const std::optional& y_routed_sf, + const std::optional& y_shared_sf, + const int& shared_sf_block_m) { + // Shifted state is all-or-nothing; FP8 uses either GEMM SF or routed and shared SF together + const bool is_shifted = shifted_prev_mix.has_value(); + DG_HOST_ASSERT(is_shifted == new_prev_mix.has_value()); + DG_HOST_ASSERT(y_routed_sf.has_value() == y_shared_sf.has_value()); + DG_HOST_ASSERT(not (y_gemm_sf.has_value() and y_routed_sf.has_value())); + DG_HOST_ASSERT(y_bf16.has_value() or y_fp8.has_value()); + DG_HOST_ASSERT(y_fp8.has_value() == (y_gemm_sf.has_value() or y_routed_sf.has_value())); + DG_HOST_ASSERT(y_shared_sf.has_value() ? shared_sf_block_m > 0 : shared_sf_block_m == 0); + + const auto [num_tokens, hidden] = get_shape<2>(x); + const auto [residual_tokens, residual_num_routes, residual_hidden] = get_shape<3>(residual); + const auto [post_mix_tokens, post_mix_num_routes, post_mix_width] = get_shape<3>(post_mix); + const auto [comb_mix_tokens, comb_mix_rows, comb_mix_cols] = get_shape<3>(comb_res_mix); + const auto [fn_num_outputs, fn_hidden] = get_shape<2>(fn); + const auto [num_mix_scales] = get_shape<1>(mix_scales); + const auto [num_mix_bases] = get_shape<1>(mix_bases); + const auto [rmsnorm_hidden] = get_shape<1>(rmsnorm_weight); + + if (num_tokens == 0) + return; + + const auto num_routes = hc_mult; + DG_HOST_ASSERT(num_routes == mhc_layout::kNumRoutes); + const auto num_hc_outputs = mhc_layout::kNumHCOutputs; + DG_HOST_ASSERT(num_tokens <= static_cast(mhc_layout::kNumMaxTokens)); + DG_HOST_ASSERT(hidden > 0 and hidden % (mhc_layout::kDefaultNumSplits * mhc_layout::BLOCK_K) == 0); + DG_HOST_ASSERT(num_sinkhorn_iters >= 1); + DG_HOST_ASSERT(num_tokens == residual_tokens and num_tokens == post_mix_tokens and num_tokens == comb_mix_tokens); + DG_HOST_ASSERT(num_routes == residual_num_routes and num_routes == post_mix_num_routes and + num_routes == comb_mix_rows and num_routes == comb_mix_cols); + DG_HOST_ASSERT(hidden == residual_hidden and post_mix_width == 1); + DG_HOST_ASSERT(num_hc_outputs == fn_num_outputs and fn_hidden == num_routes * hidden); + DG_HOST_ASSERT(num_mix_scales == 3 and num_mix_bases == num_hc_outputs); + DG_HOST_ASSERT(rmsnorm_hidden == hidden); + DG_HOST_ASSERT(new_residual.sizes() == residual.sizes()); + DG_HOST_ASSERT(new_post_mix.sizes() == post_mix.sizes()); + DG_HOST_ASSERT(new_comb_res_mix.sizes() == comb_res_mix.sizes()); + + if (is_shifted) { + const auto [prev_tokens, prev_num_routes, prev_width] = get_shape<3>(shifted_prev_mix.value()); + DG_HOST_ASSERT(prev_tokens == num_tokens and prev_num_routes == num_routes and prev_width == 1); + DG_HOST_ASSERT(new_prev_mix->sizes() == shifted_prev_mix->sizes()); + } + + const auto device = x.device(); + const auto check_dense = [&](const torch::Tensor& tensor, const torch::ScalarType& dtype) { + DG_HOST_ASSERT(tensor.scalar_type() == dtype); + DG_HOST_ASSERT(tensor.is_contiguous()); + DG_HOST_ASSERT(tensor.device() == device); + }; + + check_dense(x, torch::kBFloat16); + check_dense(residual, torch::kBFloat16); + check_dense(post_mix, torch::kFloat); + check_dense(comb_res_mix, torch::kFloat); + check_dense(fn, torch::kFloat); + check_dense(mix_scales, torch::kFloat); + check_dense(mix_bases, torch::kFloat); + check_dense(rmsnorm_weight, torch::kBFloat16); + check_dense(new_residual, torch::kBFloat16); + check_dense(new_post_mix, torch::kFloat); + check_dense(new_comb_res_mix, torch::kFloat); + if (is_shifted) { + check_dense(shifted_prev_mix.value(), torch::kFloat); + check_dense(new_prev_mix.value(), torch::kFloat); + } + + const auto check_norm_output = [&](const torch::Tensor& tensor, const torch::ScalarType& dtype) { + const auto [output_tokens, output_hidden] = get_shape<2>(tensor); + DG_HOST_ASSERT(output_tokens == num_tokens and output_hidden == hidden); + check_dense(tensor, dtype); + }; + + if (y_bf16.has_value()) + check_norm_output(y_bf16.value(), torch::kBFloat16); + if (y_fp8.has_value()) + check_norm_output(y_fp8.value(), torch::kFloat8_e4m3fn); + + const auto check_sf = [&](const torch::Tensor& tensor) { + const auto [sf_tokens, num_sf_words] = get_shape<2>(tensor); + DG_HOST_ASSERT(sf_tokens == num_tokens and + num_sf_words == hidden / static_cast(mhc_layout::kHiddenPerSFWord)); + DG_HOST_ASSERT(tensor.scalar_type() == torch::kInt); + DG_HOST_ASSERT(tensor.device() == device); + }; + + if (y_gemm_sf.has_value()) { + check_sf(y_gemm_sf.value()); + DG_HOST_ASSERT(y_gemm_sf->stride(0) == 1); + DG_HOST_ASSERT(y_gemm_sf->stride(1) >= num_tokens); + DG_HOST_ASSERT(y_gemm_sf->stride(1) % get_tma_aligned_size(1, sizeof(uint32_t)) == 0); + } + if (y_routed_sf.has_value()) { + check_sf(y_routed_sf.value()); + DG_HOST_ASSERT(y_routed_sf->is_contiguous()); + } + if (y_shared_sf.has_value()) { + check_sf(y_shared_sf.value()); + const auto num_required_rows = ceil_div(num_tokens, shared_sf_block_m) * align(shared_sf_block_m, 128); + DG_HOST_ASSERT(y_shared_sf->stride(0) == 1); + DG_HOST_ASSERT(y_shared_sf->stride(1) >= num_required_rows); + const auto num_required_elements = y_shared_sf->storage_offset() + num_required_rows + + (hidden / mhc_layout::kHiddenPerSFWord - 1) * y_shared_sf->stride(1); + const auto num_required_bytes = static_cast(num_required_elements) * y_shared_sf->element_size(); + DG_HOST_ASSERT(num_required_bytes <= y_shared_sf->storage().nbytes()); + } + + DG_HOST_ASSERT(jit->device.get_arch_major() == 10); + const auto num_sms = runtime->get_num_sms(); + const auto num_m_blocks = mhc_layout::Workspace<>::get_num_m_blocks(static_cast(num_tokens)); + const auto num_splits = get_num_splits(num_m_blocks, static_cast(hidden), static_cast(num_sms)); + const auto num_mhc_scratch_bytes = mhc_layout::Workspace<>::get_num_scratch_bytes(num_m_blocks, num_splits); + + const bool needs_bf16_scratch = not y_bf16.has_value() and is_shifted; + const auto num_bf16_scratch_bytes = needs_bf16_scratch ? static_cast(num_tokens) * hidden * x.element_size() : 0; + const auto scratch = torch::empty({static_cast(num_mhc_scratch_bytes) + num_bf16_scratch_bytes}, + x.options().dtype(torch::kByte)); + const auto& split_barriers = get_split_barriers(x.options()); + + // Shifted FP8 is derived from rounded BF16, so FP8-only calls borrow the scratch tail for that numerical boundary. + const auto y_bf16_storage = needs_bf16_scratch ? scratch.narrow(0, num_mhc_scratch_bytes, num_bf16_scratch_bytes). + view(torch::kBFloat16).view({num_tokens, hidden}) + : y_bf16.value_or(x); + + // GEMM and routed SF differ only in strides and share one primary writer. + const auto& y_primary_sf = y_routed_sf.has_value() ? y_routed_sf : y_gemm_sf; + + sm100_mega_mhc( + x, residual, post_mix, comb_res_mix, shifted_prev_mix, + fn, mix_scales, mix_bases, + hc_norm_eps, hc_pre_eps, hc_post_scale, sinkhorn_eps, num_sinkhorn_iters, + rmsnorm_weight, rmsnorm_eps, + new_residual, new_prev_mix, new_post_mix, new_comb_res_mix, + y_bf16_storage, y_bf16.has_value(), scratch, split_barriers, + y_fp8, y_primary_sf, y_shared_sf, + shared_sf_block_m, rmsnorm_scale, num_tokens, hidden, num_splits, num_sms); +} + +static void register_apis(pybind11::module_& m) { + m.def( + "mega_mhc", + &mega_mhc, + py::arg("x"), + py::arg("residual"), + py::arg("shifted_prev_mix"), + py::arg("post_mix"), + py::arg("comb_res_mix"), + py::arg("fn"), + py::arg("mix_scales"), + py::arg("mix_bases"), + py::arg("hc_mult"), + py::arg("hc_norm_eps"), + py::arg("hc_pre_eps"), + py::arg("hc_post_scale"), + py::arg("sinkhorn_eps"), + py::arg("num_sinkhorn_iters"), + py::arg("rmsnorm_weight"), + py::arg("rmsnorm_eps"), + py::arg("rmsnorm_scale"), + py::arg("new_residual"), + py::arg("new_prev_mix"), + py::arg("new_post_mix"), + py::arg("new_comb_res_mix"), + py::arg("y_bf16") = std::nullopt, + py::arg("y_fp8") = std::nullopt, + py::arg("y_gemm_sf") = std::nullopt, + py::arg("y_routed_sf") = std::nullopt, + py::arg("y_shared_sf") = std::nullopt, + py::arg("shared_sf_block_m") = 0); +} + +} // namespace deep_gemm::mega_mhc diff --git a/csrc/apis/mega.hpp b/csrc/apis/mega_moe.hpp similarity index 97% rename from csrc/apis/mega.hpp rename to csrc/apis/mega_moe.hpp index b4ab0b3f0b..309f12f040 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega_moe.hpp @@ -7,10 +7,7 @@ #include #include -#if DG_TENSORMAP_COMPATIBLE -#include "../jit/compiler.hpp" -#endif -#include "../jit/device_runtime.hpp" +#include "../runtime/runtime.hpp" #include "../jit_kernels/impls/sm100_bf16_mega_moe.hpp" #include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp" @@ -46,7 +43,7 @@ get_symm_buffer_size_for_mega_moe( // Ring capacity: worst-case live pool blocks over all candidate BLOCK_M; mirrors the kernel assert. // TODO: we temporarily assume the SM count is consistent with the runtime value - const auto num_sms = device_runtime->get_num_sms(); + const auto num_sms = runtime->get_num_sms(); const auto num_experts_per_rank = num_experts / num_ranks; const auto num_active_topk = std::min(num_topk, num_experts_per_rank); const auto num_max_routed_tokens = num_max_tokens_per_rank * num_ranks * num_active_topk; @@ -188,13 +185,14 @@ static void fp8_fp4_mega_moe( // Tensor checks 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); - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); const auto [num_experts_per_rank, intermediate_hidden_2, hidden] = check_grouped_ab_fp8_fp4(l1_weights, cute::UMMA::Major::K, arch_major); const auto [num_experts_per_rank_, hidden_, intermediate_hidden] = check_grouped_ab_fp8_fp4(l2_weights, cute::UMMA::Major::K, arch_major); - DG_HOST_ASSERT(l1_weights.scalar_type() == kPackedFP4); - DG_HOST_ASSERT(l2_weights.scalar_type() == kPackedFP4); + const auto weight_dtype = l1_weights.scalar_type(); + DG_HOST_ASSERT(weight_dtype == torch::kFloat8_e4m3fn or weight_dtype == kPackedFP4); + DG_HOST_ASSERT(l2_weights.scalar_type() == weight_dtype); 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_); @@ -246,7 +244,8 @@ static void fp8_fp4_mega_moe( num_ranks, num_experts, num_max_tokens_per_rank, num_topk, hidden, intermediate_hidden, - "fp8xfp4", activation, num_shared_experts + weight_dtype == torch::kFloat8_e4m3fn ? "fp8xfp8" : "fp8xfp4", + activation, num_shared_experts ); DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast(num_required_bytes)); DG_HOST_ASSERT(num_experts == num_experts_); @@ -281,7 +280,7 @@ static void fp8_fp4_mega_moe( // Zero the entire symmetric buffer for debug mode // NOTES: caller must re-copy inputs into the buffer before each kernel call - if (get_env("DG_COMM_KERNEL_DEBUG")) + if (deep_jit::get_env("DG_COMM_KERNEL_DEBUG")) sym_buffer.zero_(); } @@ -313,7 +312,7 @@ static void bf16_mega_moe( // Tensor checks 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); - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); 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(l1_weights.scalar_type() == torch::kBFloat16); @@ -389,18 +388,16 @@ static void bf16_mega_moe( // Zero the entire symmetric buffer for debug mode // NOTES: caller must re-copy inputs into the buffer before each kernel call - if (get_env("DG_COMM_KERNEL_DEBUG")) + if (deep_jit::get_env("DG_COMM_KERNEL_DEBUG")) sym_buffer.zero_(); } 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); m.def("get_block_m_for_mega_moe", &get_block_m_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("bf16_mega_moe", &bf16_mega_moe); -#endif } } // namespace deep_gemm::mega diff --git a/csrc/indexing/main.cu b/csrc/indexing/main.cu index 2c67e66923..0c05ac7684 100644 --- a/csrc/indexing/main.cu +++ b/csrc/indexing/main.cu @@ -20,11 +20,11 @@ // Layout kernels #include -#include // Mega kernels #include #include +#include using namespace deep_gemm; diff --git a/csrc/jit/cache.hpp b/csrc/jit/cache.hpp deleted file mode 100644 index ddc763d08d..0000000000 --- a/csrc/jit/cache.hpp +++ /dev/null @@ -1,31 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "kernel_runtime.hpp" - -namespace deep_gemm { - -class KernelRuntimeCache { - std::unordered_map> cache; - -public: - // TODO: consider cache capacity - KernelRuntimeCache() = default; - - std::shared_ptr get(const std::filesystem::path& dir_path) { - // Hit the runtime cache - if (const auto iterator = cache.find(dir_path); iterator != cache.end()) - return iterator->second; - - if (KernelRuntime::check_validity(dir_path)) - return cache[dir_path] = std::make_shared(dir_path); - return nullptr; - } -}; - -static auto kernel_runtime_cache = std::make_shared(); - -} // namespace deep_gemm diff --git a/csrc/jit/compiler.hpp b/csrc/jit/compiler.hpp deleted file mode 100644 index 7d85a5f556..0000000000 --- a/csrc/jit/compiler.hpp +++ /dev/null @@ -1,362 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include -#include - -#include "../utils/exception.hpp" -#include "../utils/format.hpp" -#include "../utils/hash.hpp" -#include "../utils/lazy_init.hpp" -#include "../utils/system.hpp" -#include "cache.hpp" -#include "device_runtime.hpp" -#include "include_parser.hpp" - -namespace deep_gemm { - -class Compiler { -public: - static std::filesystem::path library_root_path; - static std::filesystem::path library_include_path; - static std::filesystem::path cuda_home; - static std::filesystem::path cuobjdump_path; - - static void prepare_init(const std::string& library_root_path, - const std::string& cuda_home_path_by_python) { - Compiler::library_root_path = library_root_path; - Compiler::library_include_path = Compiler::library_root_path / "include"; - Compiler::cuda_home = cuda_home_path_by_python; - Compiler::cuobjdump_path = Compiler::cuda_home / "bin" / "cuobjdump"; - } - - std::string signature, flags; - std::filesystem::path cache_dir_path; - - Compiler() { - // Check `prepare_init` - DG_HOST_ASSERT(not library_root_path.empty()); - DG_HOST_ASSERT(not library_include_path.empty()); - DG_HOST_ASSERT(not cuda_home.empty()); - DG_HOST_ASSERT(not cuobjdump_path.empty()); - - // Cache settings - cache_dir_path = std::filesystem::path(get_env("HOME")) / ".deep_gemm"; - if (const auto env_cache_dir_path = get_env("DG_JIT_CACHE_DIR"); not env_cache_dir_path.empty()) - cache_dir_path = env_cache_dir_path; - - // The compiler flags applied to all derived compilers - signature = "unknown-compiler"; - flags = fmt::format("-std=c++{} --diag-suppress=39,161,174,177,186,940 " - "--ptxas-options=--register-usage-level=10", - get_env("DG_JIT_CPP_STANDARD", 20)); - if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0) or get_env("DG_JIT_PTXAS_CHECK", 0)) - flags += " --ptxas-options=--verbose,--warn-on-local-memory-usage"; - if (get_env("DG_JIT_WITH_LINEINFO", 0)) - flags += " -Xcompiler -rdynamic -lineinfo"; - } - - virtual ~Compiler() = default; - - std::filesystem::path make_tmp_dir() const { - return make_dirs(cache_dir_path / "tmp"); - } - - static void fsync_path(const std::filesystem::path& path) { - const auto fd = ::open(path.c_str(), O_RDONLY); - if (fd >= 0) { - ::fsync(fd); - ::close(fd); - } - } - - // Recursively fsync a directory: files and subdirectories first (bottom-up), then the directory itself - // NOTES: ensures data and directory entries are visible on other nodes in distributed filesystems - static void fsync_dir(const std::filesystem::path& dir_path) { // NOLINT(*-no-recursion) - for (const auto& entry: std::filesystem::directory_iterator(dir_path)) { - if (entry.is_directory()) - fsync_dir(entry.path()); - else if (entry.is_regular_file()) - fsync_path(entry.path()); - } - fsync_path(dir_path); - } - - static void put(const std::filesystem::path& path, const std::string& data) { - std::ofstream out(path, std::ios::binary); - DG_HOST_ASSERT(out.write(data.data(), data.size())); - out.close(); - - // NOTES: fsync to ensure the data is visible to other processes (e.g., NVCC) - // on distributed filesystems, where `close()` alone does not guarantee persistence - fsync_path(path); - } - - std::shared_ptr build(const std::string& name, const std::string& code) const { - const auto kernel_signature = fmt::format("{}$${}$${}$${}", name, signature, flags, code); - const auto dir_path = cache_dir_path / "cache" / fmt::format("kernel.{}.{}", name, get_hex_digest(kernel_signature)); - - // Hit the runtime cache - if (const auto runtime = kernel_runtime_cache->get(dir_path); runtime != nullptr) - return runtime; - - // Compile into a temporary directory, then atomically rename the whole directory - // NOTES: renaming a directory is atomic on both local and distributed filesystems, - // avoiding the stale inode issue that occurs when renaming individual files - const auto tmp_dir_path = make_tmp_dir() / get_uuid(); - make_dirs(tmp_dir_path); - - // Compile into the temporary directory - const auto tmp_cubin_path = tmp_dir_path / "kernel.cubin"; - if (get_env("DG_JIT_DUMP_ASM") or get_env("DG_JIT_DUMP_PTX")) { - const auto tmp_ptx_path = tmp_dir_path / "kernel.ptx"; - compile(code, tmp_dir_path, tmp_cubin_path, tmp_ptx_path); - } else { - compile(code, tmp_dir_path, tmp_cubin_path); - } - - // Disassemble if needed - if (get_env("DG_JIT_DUMP_ASM") or get_env("DG_JIT_DUMP_SASS")) { - const auto tmp_sass_path = tmp_dir_path / "kernel.sass"; - disassemble(tmp_cubin_path, tmp_sass_path); - } - - // Fsync before rename to ensure visibility on distributed filesystems - fsync_dir(tmp_dir_path); - - // Atomically rename the temporary directory to the final cache path - // NOTES: if another rank already created dir_path, rename will fail — that's fine - make_dirs(dir_path.parent_path()); - std::error_code error_code; - std::filesystem::rename(tmp_dir_path, dir_path, error_code); - if (error_code) { - // Another rank beat us, then clean up our dir and use the existing one - // NOTES: avoid `std::filesystem::remove_all` here — it can segfault on - // distributed filesystems, when concurrent processes operate - // on the same parent directory, causing stale directory entries - safe_remove_all(tmp_dir_path); - } - - // Put into the runtime cache - const auto runtime = kernel_runtime_cache->get(dir_path); - DG_HOST_ASSERT(runtime != nullptr); - return runtime; - } - - static void disassemble(const std::filesystem::path &cubin_path, const std::filesystem::path &sass_path) { - // Disassemble the CUBIN file to SASS - const auto command = fmt::format("{} --dump-sass {} > {}", cuobjdump_path.c_str(), cubin_path.c_str(), sass_path.c_str()); - if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PRINT_COMPILER_COMMAND", 0)) - printf("Running cuobjdump command: %s\n", command.c_str()); - const auto [return_code, output] = call_external_command(command); - if (return_code != 0) { - printf("cuobjdump failed: %s\n", output.c_str()); - DG_HOST_ASSERT(false and "cuobjdump failed"); - } - } - - virtual void compile(const std::string &code, const std::filesystem::path& dir_path, const std::filesystem::path &cubin_path, const std::optional &ptx_path = std::nullopt) const = 0; -}; - -DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_root_path); -DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, library_include_path); -DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, cuda_home); -DG_DECLARE_STATIC_VAR_IN_CLASS(Compiler, cuobjdump_path); - -class NVCCCompiler final: public Compiler { - std::filesystem::path nvcc_path; - - std::pair get_nvcc_version() const { - DG_HOST_ASSERT(std::filesystem::exists(nvcc_path)); - - // Call the version command - const auto command = std::string(nvcc_path) + " --version"; - const auto [return_code, output] = call_external_command(command); - DG_HOST_ASSERT(return_code == 0); - - // The version should be at least 12.3, for the best performance with 12.9 - int major, minor; - std::smatch match; - DG_HOST_ASSERT(std::regex_search(output, match, std::regex(R"(release (\d+\.\d+))"))); - std::sscanf(match[1].str().c_str(), "%d.%d", &major, &minor); - DG_HOST_ASSERT((major > 12 or (major == 12 and minor >= 3)) and "NVCC version should be >= 12.3"); - if (major == 12 and minor < 9) - printf("Warning: please use at least NVCC 12.9 for the best DeepGEMM performance\n"); - return {major, minor}; - } - -public: - NVCCCompiler() { - // Override the compiler signature - nvcc_path = cuda_home / "bin" / "nvcc"; - if (const auto env_nvcc_path = get_env("DG_JIT_NVCC_COMPILER"); not env_nvcc_path.empty()) - nvcc_path = env_nvcc_path; - const auto [nvcc_major, nvcc_minor] = get_nvcc_version(); - signature = fmt::format("NVCC{}.{}", nvcc_major, nvcc_minor); - - // The override the compiler flags - // Only NVCC >= 12.9 supports arch-specific family suffix - const auto arch = device_runtime->get_arch(false, nvcc_major > 12 or nvcc_minor >= 9); - flags = fmt::format("{} -I{} --gpu-architecture=sm_{} " - "--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi " - "-O3 --expt-relaxed-constexpr --expt-extended-lambda", - flags, library_include_path.c_str(), arch); - } - - void compile(const std::string &code, const std::filesystem::path& dir_path, - const std::filesystem::path &cubin_path, - const std::optional &ptx_path) const override { - // Write the code into the cache directory - const auto code_path = dir_path / "kernel.cu"; - put(code_path, code); - - // Compile - // Avoid cwd files shadowing C++ standard library headers - const auto compile_dir = make_tmp_dir(); - const auto command = fmt::format("cd {} && {} {} -cubin -o {} {}", - compile_dir.c_str(), nvcc_path.c_str(), code_path.c_str(), cubin_path.c_str(), flags); - if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PRINT_COMPILER_COMMAND", 0)) - printf("Running NVCC command: %s\n", command.c_str()); - const auto [return_code, output] = call_external_command(command); - if (return_code != 0) { - printf("NVCC compilation failed: %s\n", output.c_str()); - DG_HOST_ASSERT(false and "NVCC compilation failed"); - } - - // Compile to PTX if needed - if (ptx_path.has_value()) { - const auto ptx_command = fmt::format("cd {} && {} {} -ptx -o {} {}", - compile_dir.c_str(), nvcc_path.c_str(), code_path.c_str(), ptx_path->c_str(), flags); - if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PRINT_COMPILER_COMMAND", 0)) - printf("Running NVCC PTX command: %s\n", ptx_command.c_str()); - const auto [ptx_return_code, ptx_output] = call_external_command(ptx_command); - if (ptx_return_code != 0) { - printf("NVCC PTX compilation failed: %s\n", ptx_output.c_str()); - DG_HOST_ASSERT(false and "NVCC PTX compilation failed"); - } - } - - // Check local memory usage - if (get_env("DG_JIT_PTXAS_CHECK", 0)) - DG_HOST_ASSERT(not std::regex_search(output, std::regex(R"(Local memory used)"))); - - // Print PTXAS log - if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0)) - printf("%s", output.c_str()); - } -}; - -class NVRTCCompiler final: public Compiler { -public: - NVRTCCompiler() { - // Override the compiler signature - int major, minor; - DG_NVRTC_CHECK(nvrtcVersion(&major, &minor)); - signature = fmt::format("NVRTC{}.{}", major, minor); - DG_HOST_ASSERT((major > 12 or (major == 12 and minor >= 3)) and "NVRTC version should be >= 12.3"); - - // Build include directories list - std::string include_dirs; - include_dirs += fmt::format("-I{} ", library_include_path.string()); - include_dirs += fmt::format("-I{} ", (cuda_home / "include").string()); - - // Add PCH support for version 12.8 and above - // NOTES: PCH is vital for compilation speed - std::string pch_flags; - if (major > 12 or minor >= 8) { - pch_flags = "--pch "; - if (get_env("DG_JIT_DEBUG", 0)) - pch_flags += "--pch-verbose=true "; - } - - // Override the compiler flags - // Only NVRTC >= 12.9 supports arch-specific family suffix - const auto arch = device_runtime->get_arch(false, major > 12 or minor >= 9); - flags = fmt::format("{} {}--gpu-architecture=sm_{} -default-device {} --device-int128", - flags, include_dirs, arch, pch_flags); - } - - void compile(const std::string &code, const std::filesystem::path& dir_path, - const std::filesystem::path &cubin_path, - const std::optional &ptx_path) const override { - // Write the code into the cache directory - const auto code_path = dir_path / "kernel.cu"; - put(code_path, code); - - // Parse compilation options - std::istringstream iss(flags); - std::vector options; - std::string option; - while (iss >> option) - options.push_back(option); - - // Convert to C-style string array for NVRTC - std::vector option_cstrs; - for (const auto& opt: options) - option_cstrs.push_back(opt.c_str()); - - // Print compiler command if requested - if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PRINT_COMPILER_COMMAND", 0)) { - printf("Compiling JIT runtime with NVRTC options: "); - for (const auto& opt: options) - printf("%s ", opt.c_str()); - printf("\n"); - } - - // Create NVRTC program and compile - nvrtcProgram program; - DG_NVRTC_CHECK(nvrtcCreateProgram(&program, code.c_str(), "kernel.cu", 0, nullptr, nullptr)); - const auto compile_result = nvrtcCompileProgram(program, static_cast(option_cstrs.size()), option_cstrs.data()); - - // Get and print compiler log - size_t log_size; - DG_NVRTC_CHECK(nvrtcGetProgramLogSize(program, &log_size)); - if (get_env("DG_JIT_DEBUG", 0) or compile_result != NVRTC_SUCCESS) { - if (compile_result != NVRTC_SUCCESS) - DG_HOST_ASSERT(log_size > 1); - if (log_size > 1) { - std::string compilation_log(log_size, '\0'); - DG_NVRTC_CHECK(nvrtcGetProgramLog(program, compilation_log.data())); - printf("NVRTC log: %s\n", compilation_log.c_str()); - } - } - - if (ptx_path.has_value()) { - // Get PTX size and data if needed - size_t ptx_size; - DG_NVRTC_CHECK(nvrtcGetPTXSize(program, &ptx_size)); - std::string ptx_data(ptx_size, '\0'); - DG_NVRTC_CHECK(nvrtcGetPTX(program, ptx_data.data())); - - // Write into the file system - put(ptx_path.value(), ptx_data); - } - - // Get CUBIN size and data - size_t cubin_size; - DG_NVRTC_CHECK(nvrtcGetCUBINSize(program, &cubin_size)); - std::string cubin_data(cubin_size, '\0'); - DG_NVRTC_CHECK(nvrtcGetCUBIN(program, cubin_data.data())); - - // Write into the file system - put(cubin_path, cubin_data); - - // Cleanup - DG_NVRTC_CHECK(nvrtcDestroyProgram(&program)); - } -}; - -static auto compiler = LazyInit([]() -> std::shared_ptr { - if (get_env("DG_JIT_USE_NVRTC", 0)) { - return std::make_shared(); - } else { - return std::make_shared(); - } -}); - -} // namespace deep_gemm diff --git a/csrc/jit/device_runtime.hpp b/csrc/jit/device_runtime.hpp deleted file mode 100644 index 2321aded7a..0000000000 --- a/csrc/jit/device_runtime.hpp +++ /dev/null @@ -1,138 +0,0 @@ -#pragma once - -#include -#include -#include - -#include "../utils/exception.hpp" -#include "../utils/lazy_init.hpp" - -#define PYTORCH_SUPPORTS_GET_CUBLASLT_HANDLE (TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 3)) - -namespace deep_gemm { - -class DeviceRuntime { - int num_sms = 0, tc_util = 0; - bool enable_pdl = false; - std::shared_ptr cached_prop; - - // cuBLASLt utils - static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024; - -public: - // Create the cuBLASLt handle ourselves - cublasLtHandle_t cublaslt_handle; - torch::Tensor cublaslt_workspace; - bool use_pytorch_managed_cublaslt_handle; - bool use_temp_cublaslt_workspace; - - explicit DeviceRuntime() { - - // Whether to use PyTorch cuBLASLt - // By default, we don't use it, - // as `at::cuda::getCurrentCUDABlasLtHandle` has large CPU overhead with some PyTorch versions - use_pytorch_managed_cublaslt_handle = get_env("DG_USE_PYTORCH_CUBLASLT_HANDLE", 0) > 0; -#if not PYTORCH_SUPPORTS_GET_CUBLASLT_HANDLE - DG_HOST_ASSERT(not use_pytorch_managed_cublaslt_handle and "PyTorch does not support to get cuBLASLt handle"); -#endif - - // Whether to create workspace tensor on each call instead of holding one. - // Enabled by compute-sanitizer tests, which trigger `cudaErrorCudartUnloading` - // when the workspace tensor is destructed after CUDA driver shutdown. - use_temp_cublaslt_workspace = get_env("DG_USE_TEMP_CUBLASLT_WORKSPACE", 0) > 0; - - if (not use_pytorch_managed_cublaslt_handle) - DG_CUBLASLT_CHECK(cublasLtCreate(&cublaslt_handle)); - - if (not use_temp_cublaslt_workspace) - cublaslt_workspace = torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA)); - } - - ~DeviceRuntime() noexcept(false) { - if (not use_pytorch_managed_cublaslt_handle) - DG_CUBLASLT_CHECK(cublasLtDestroy(cublaslt_handle)); - } - - cublasLtHandle_t get_cublaslt_handle() const { -#if PYTORCH_SUPPORTS_GET_CUBLASLT_HANDLE - if (use_pytorch_managed_cublaslt_handle) - return at::cuda::getCurrentCUDABlasLtHandle(); -#endif - - // Self-managed handle - return cublaslt_handle; - } - - torch::Tensor get_cublaslt_workspace() const { - if (use_temp_cublaslt_workspace) - return torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA)); - return cublaslt_workspace; - } - - std::shared_ptr get_prop() { - if (cached_prop == nullptr) { - int device_idx; - cudaDeviceProp prop; - DG_CUDA_RUNTIME_CHECK(cudaGetDevice(&device_idx)); - DG_CUDA_RUNTIME_CHECK(cudaGetDeviceProperties(&prop, device_idx)); - cached_prop = std::make_shared(prop); - } - return cached_prop; - } - - std::pair get_arch_pair() { - const auto prop = get_prop(); - return {prop->major, prop->minor}; - } - - std::string get_arch(const bool& number_only = false, - const bool& support_arch_family = false) { - const auto [major, minor] = get_arch_pair(); - if (major == 10 and minor != 1) { - if (number_only) - return "100"; - return support_arch_family ? "100f" : "100a"; - } - return std::to_string(major * 10 + minor) + (number_only ? "" : "a"); - } - - int get_arch_major() { - return get_arch_pair().first; - } - - void set_num_sms(const int& new_num_sms) { - DG_HOST_ASSERT(0 <= new_num_sms and new_num_sms <= get_prop()->multiProcessorCount); - num_sms = new_num_sms; - } - - int get_num_sms() { - if (num_sms == 0) - num_sms = get_prop()->multiProcessorCount; - return num_sms; - } - - int get_l2_cache_size() { - return get_prop()->l2CacheSize; - } - - void set_tc_util(const int& new_tc_util) { - DG_HOST_ASSERT(0 <= new_tc_util and new_tc_util <= 100); - tc_util = new_tc_util; - } - - int get_tc_util() const { - return tc_util == 0 ? 100 : tc_util; - } - - void set_pdl(const bool& new_enable_pdl) { - enable_pdl = new_enable_pdl; - } - - bool get_pdl() const { - return enable_pdl; - } -}; - -static auto device_runtime = LazyInit([](){ return std::make_shared(); }); - -} // namespace deep_gemm diff --git a/csrc/jit/handle.hpp b/csrc/jit/handle.hpp deleted file mode 100644 index be3bc31c07..0000000000 --- a/csrc/jit/handle.hpp +++ /dev/null @@ -1,222 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -#include "../utils/exception.hpp" -#include "../utils/compatibility.hpp" - -namespace deep_gemm { - -// Lazy loading all driver symbols -static void* get_driver_handle() { - static void* handle = nullptr; - if (handle == nullptr) { - handle = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL); - DG_HOST_ASSERT(handle != nullptr and "Failed to load CUDA driver `libcuda.so.1`"); - } - return handle; -} - -// Macro to define wrapper functions named `lazy_cu{API name}` -#define DECL_LAZY_CUDA_DRIVER_FUNCTION(name) \ -template \ -static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \ - using FuncType = decltype(&(name)); \ - static FuncType func = nullptr; \ - if (func == nullptr) { \ - func = reinterpret_cast(dlsym(get_driver_handle(), #name)); \ - DG_HOST_ASSERT(func != nullptr and "Failed to load CUDA driver API"); \ - } \ - return func(std::forward(args)...); \ -} - -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorName); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorString); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuFuncSetAttribute); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleLoad); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleUnload); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleGetFunction); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLibraryLoadFromFile); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLibraryUnload); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuKernelGetFunction); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLaunchKernelEx); -DECL_LAZY_CUDA_DRIVER_FUNCTION(cuTensorMapEncodeTiled); - -#if CUDART_VERSION >= 12080 and defined(DG_JIT_USE_RUNTIME_API) - -// Use CUDA runtime API -using LibraryHandle = cudaLibrary_t; -using KernelHandle = cudaKernel_t; -using LaunchConfigHandle = cudaLaunchConfig_t; -using LaunchAttrHandle = cudaLaunchAttribute; - -#define DG_CUDA_UNIFIED_CHECK DG_CUDA_RUNTIME_CHECK - -static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const std::string& func_name, - LibraryHandle *library_opt = nullptr) { - LibraryHandle library; - KernelHandle kernel{}; - DG_CUDA_RUNTIME_CHECK(cudaLibraryLoadFromFile(&library, cubin_path.c_str(), nullptr, nullptr, 0, nullptr, nullptr, 0)); - DG_CUDA_RUNTIME_CHECK(cudaLibraryGetKernel(&kernel, library, func_name.c_str())); - - if (library_opt != nullptr) - *library_opt = library; - return kernel; -} - -static void unload_library(const LibraryHandle& library) { - const auto error = cudaLibraryUnload(library); - DG_HOST_ASSERT(error == cudaSuccess or error == cudaErrorCudartUnloading); -} - -static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, - const cudaStream_t& stream, const int& smem_size, - const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, const bool& enable_pdl) { - if (smem_size > 0) - DG_CUDA_RUNTIME_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size)); - - LaunchConfigHandle config; - config.gridDim = grid_dim; - config.blockDim = block_dim; - config.dynamicSmemBytes = smem_size; - config.stream = stream; - - // Create attributes - // NOTES: must use `static` or the `attr` will be deconstructed - static LaunchAttrHandle attrs[2]; - config.numAttrs = 0; - config.attrs = attrs; - - // Cluster size - if (cluster_dim > 1) { - auto& attr = attrs[config.numAttrs ++]; - attr.id = cudaLaunchAttributeClusterDimension; - attr.val.clusterDim = {static_cast(cluster_dim), 1, 1}; - } - - // Dependent kernel launch - if (enable_pdl) { - auto& attr = attrs[config.numAttrs ++]; - attr.id = cudaLaunchAttributeProgrammaticStreamSerialization; - attr.val.programmaticStreamSerializationAllowed = 1; - } - - return config; -} - -template -static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) { - void *ptr_args[] = { &args... }; - return cudaLaunchKernelExC(&config, kernel, ptr_args); -} - -#else - -// Use CUDA driver API -using KernelHandle = CUfunction; -using LaunchConfigHandle = CUlaunchConfig; -using LaunchAttrHandle = CUlaunchAttribute; - -// `cuLibraryEnumerateKernels` is supported since CUDA Driver API 12.4 -#if CUDA_VERSION >= 12040 - #define DG_JIT_USE_LIBRARY_ENUM_KERNELS - DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLibraryGetKernelCount); - DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLibraryEnumerateKernels); - using LibraryHandle = CUlibrary; -#else - using LibraryHandle = CUmodule; -#endif - -#define DG_CUDA_UNIFIED_CHECK DG_CUDA_DRIVER_CHECK - -static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const std::string& func_name, - LibraryHandle *library_opt = nullptr) { - LibraryHandle library; - KernelHandle kernel; - -#ifdef DG_JIT_USE_LIBRARY_ENUM_KERNELS - DG_CUDA_DRIVER_CHECK(lazy_cuLibraryLoadFromFile(&library, cubin_path.c_str(), nullptr, nullptr, 0, nullptr, nullptr, 0)); - unsigned int num_kernels; - DG_CUDA_DRIVER_CHECK(lazy_cuLibraryGetKernelCount(&num_kernels, library)); - if (num_kernels != 1) { - const auto dir_path = cubin_path.parent_path(); - printf("Corrupted JIT cache directory (expected 1 kernel, found %u): %s, " - "please run `rm -rf %s` and restart your task.\n", - num_kernels, dir_path.c_str(), dir_path.c_str()); - DG_HOST_ASSERT(false and "Corrupted JIT cache directory"); - } - - CUkernel cu_kernel; - DG_CUDA_DRIVER_CHECK(lazy_cuLibraryEnumerateKernels(&cu_kernel, 1, library)); - DG_CUDA_DRIVER_CHECK(lazy_cuKernelGetFunction(&kernel, cu_kernel)); -#else - DG_CUDA_DRIVER_CHECK(lazy_cuModuleLoad(&library, cubin_path.c_str())); - DG_CUDA_DRIVER_CHECK(lazy_cuModuleGetFunction(&kernel, library, func_name.c_str())); -#endif - - if (library_opt != nullptr) - *library_opt = library; - return kernel; -} - -static void unload_library(const LibraryHandle& library) { -#ifdef DG_JIT_USE_LIBRARY_ENUM_KERNELS - const auto error = lazy_cuLibraryUnload(library); -#else - const auto error = lazy_cuModuleUnload(library); -#endif - DG_HOST_ASSERT(error == CUDA_SUCCESS or error == CUDA_ERROR_DEINITIALIZED); -} - -static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel, - const cudaStream_t& stream, const int& smem_size, - const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim, const bool& enable_pdl) { - if (smem_size > 0) - DG_CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size)); - - LaunchConfigHandle config; - config.gridDimX = grid_dim.x; - config.gridDimY = grid_dim.y; - config.gridDimZ = grid_dim.z; - config.blockDimX = block_dim.x; - config.blockDimY = block_dim.y; - config.blockDimZ = block_dim.z; - config.sharedMemBytes = smem_size; - config.hStream = stream; - - // Create attributes - // NOTES: must use `static` or the `attr` will be deconstructed - static LaunchAttrHandle attrs[2]; - config.numAttrs = 0; - config.attrs = attrs; - - // Cluster size - if (cluster_dim > 1) { - auto& attr = attrs[config.numAttrs ++]; - attr.id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; - attr.value.clusterDim.x = static_cast(cluster_dim); - attr.value.clusterDim.y = 1; - attr.value.clusterDim.z = 1; - } - - // Dependent kernel launch - if (enable_pdl) { - auto& attr = attrs[config.numAttrs ++]; - attr.id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; - attr.value.programmaticStreamSerializationAllowed = 1; - } - - return config; -} - -template -static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) { - void *ptr_args[] = { &args... }; - return lazy_cuLaunchKernelEx(&config, kernel, ptr_args, nullptr); -} -#endif - -} // namespace deep_gemm diff --git a/csrc/jit/include_parser.hpp b/csrc/jit/include_parser.hpp deleted file mode 100644 index 99f2663c78..0000000000 --- a/csrc/jit/include_parser.hpp +++ /dev/null @@ -1,80 +0,0 @@ -#pragma once - -#include -#include -#include -#include - -#include "../utils/format.hpp" -#include "../utils/system.hpp" - -namespace deep_gemm { - -class IncludeParser { - std::unordered_map> cache; - - static std::vector get_includes(const std::string& code, const std::filesystem::path& file_path = "") { - std::vector includes; - const std::regex pattern(R"(#\s*include\s*[<"][^>"]+[>"])"); - std::sregex_iterator iter(code.begin(), code.end(), pattern); - const std::sregex_iterator end; - - // TODO: parse relative paths as well - for (; iter != end; ++ iter) { - const auto include_str = iter->str(); - const int len = include_str.length(); - if (include_str.substr(0, 10) == "#include <" and include_str[len - 1] == '>' and include_str[10] != ' ' and include_str[len - 2] != ' ') { - std::string filename = include_str.substr(10, len - 11); - if (filename.substr(0, 9) == "deep_gemm") // We only parse `` - includes.push_back(filename); - } else { - std::string error_info = fmt::format("Non-standard include: {}", include_str); - if (file_path != "") - error_info += fmt::format(" ({})", file_path.string()); - DG_HOST_UNREACHABLE(error_info); - } - } - return includes; - } - -public: - static std::filesystem::path library_include_path; - - static void prepare_init(const std::string& library_root_path) { - library_include_path = std::filesystem::path(library_root_path) / "include"; - } - - std::string get_hash_value(const std::string& code, const bool& exclude_code = true) { - std::stringstream ss; - for (const auto& i: get_includes(code)) - ss << get_hash_value_by_path(library_include_path / i) << "$"; - if (not exclude_code) - ss << "#" << get_hex_digest(code); - return get_hex_digest(ss.str()); - } - - std::string get_hash_value_by_path(const std::filesystem::path& path) { - // Check whether hit in cache - // ReSharper disable once CppUseAssociativeContains - if (cache.count(path) > 0) { - const auto opt = cache[path]; - if (not opt.has_value()) - DG_HOST_UNREACHABLE(fmt::format("Circular include may occur: {}", path.string())); - return opt.value(); - } - - // Read file and calculate hash recursively - std::ifstream in(path); - if (not in.is_open()) - DG_HOST_UNREACHABLE(fmt::format("Failed to open: {}", path.string())); - std::string code((std::istreambuf_iterator(in)), std::istreambuf_iterator()); - cache[path] = std::nullopt; - return (cache[path] = get_hash_value(code, false)).value(); - } -}; - -DG_DECLARE_STATIC_VAR_IN_CLASS(IncludeParser, library_include_path); - -static auto include_parser = std::make_shared(); - -} // namespace deep_gemm diff --git a/csrc/jit/kernel_runtime.hpp b/csrc/jit/kernel_runtime.hpp deleted file mode 100644 index 40597fb448..0000000000 --- a/csrc/jit/kernel_runtime.hpp +++ /dev/null @@ -1,165 +0,0 @@ -#pragma once - -#include - -#include "../utils/exception.hpp" -#include "../utils/format.hpp" -#include "../utils/system.hpp" -#include "device_runtime.hpp" -#include "handle.hpp" -#include "include_parser.hpp" - -namespace deep_gemm { - -struct LaunchArgs { - std::pair grid_dim; - int num_threads; - int smem_size; - int cluster_dim; - bool enable_pdl; - - LaunchArgs(const int& grid_dim_x, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& enable_pdl = true): - grid_dim({grid_dim_x, 1}), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), enable_pdl(enable_pdl) {} - - LaunchArgs(const std::pair& grid_dim, const int& num_threads, const int& smem_size = 0, const int& cluster_dim = 1, const bool& enable_pdl = true): - grid_dim(grid_dim), num_threads(num_threads), smem_size(smem_size), cluster_dim(cluster_dim), enable_pdl(enable_pdl) {} -}; - -class KernelRuntime final { -public: - static std::filesystem::path cuda_home; - - LibraryHandle library; - KernelHandle kernel; - - explicit KernelRuntime(const std::filesystem::path& dir_path) { - // Check `prepare_init` - DG_HOST_ASSERT(not cuda_home.empty()); - - // NOLINT(*-pro-type-member-init) - const auto cuobjdump_path = cuda_home / "bin" / "cuobjdump"; - const auto cubin_path = dir_path / "kernel.cubin"; - if (get_env("DG_JIT_DEBUG")) - printf("Loading CUBIN: %s\n", cubin_path.c_str()); - - // Record start time - std::chrono::high_resolution_clock::time_point start_time; - if (get_env("DG_JIT_DEBUG") or get_env("DG_JIT_PRINT_LOAD_TIME")) - start_time = std::chrono::high_resolution_clock::now(); - -#ifdef DG_JIT_USE_LIBRARY_ENUM_KERNELS - // Load from the library - kernel = load_kernel(cubin_path, {}, &library); -#else - // Find the only symbol - // TODO: use kernel enumeration for newer drivers - const std::vector illegal_names = {"vprintf", "__instantiate_kernel", "__internal", "__assertfail"}; - const auto [exit_code, symbols] = call_external_command(fmt::format("{} -symbols {}", cuobjdump_path.c_str(), cubin_path.c_str())); - DG_HOST_ASSERT(exit_code == 0); - std::istringstream iss(symbols); - std::vector symbol_names; - for (std::string line; std::getline(iss, line); ) { - if (line.find("STT_FUNC") == 0 and line.find("STO_ENTRY") != std::string::npos and - std::none_of(illegal_names.begin(), illegal_names.end(), - [&](const auto name) { return line.find(name) != std::string::npos; })) { - const auto last_space = line.rfind(' '); - symbol_names.push_back(line.substr(last_space + 1)); - } - } - - // Print symbols - if (symbol_names.size() != 1 or get_env("DG_JIT_DEBUG")) { - printf("Symbols: "); - printf(" > CUBIN: %s\n", cubin_path.c_str()); - printf(" > Raw symbols: %s\n", symbols.c_str()); - printf(" > Parsed symbols:\n"); - for (const auto& symbol: symbol_names) - printf(" > %s, ", symbol.c_str()); - } - DG_HOST_ASSERT(symbol_names.size() == 1); - - // Load from the library - kernel = load_kernel(cubin_path, symbol_names[0], &library); -#endif - - // Print load time - if (get_env("DG_JIT_DEBUG") or get_env("DG_JIT_PRINT_LOAD_TIME")) { - std::chrono::duration load_time = std::chrono::high_resolution_clock::now() - start_time; - printf("Load time (%s): %.2lf ms\n", dir_path.c_str(), load_time.count()); - } - } - - static void prepare_init(const std::string& cuda_home_path_by_python) { - cuda_home = cuda_home_path_by_python; - } - - static bool check_validity(const std::filesystem::path& dir_path) { - if (not std::filesystem::exists(dir_path)) - return false; - - // NOTES: if the directory exists, `kernel.cu` and `kernel.cubin` must both exist, - // because the directory is created atomically via rename - if (not std::filesystem::exists(dir_path / "kernel.cu") or - not std::filesystem::exists(dir_path / "kernel.cubin")) { - printf("Corrupted JIT cache directory (missing kernel.cu or kernel.cubin): %s, " - "please run `rm -rf %s` and restart your task.\n", - dir_path.c_str(), dir_path.c_str()); - DG_HOST_ASSERT(false and "Corrupted JIT cache directory"); - } - return true; - } - - ~KernelRuntime() noexcept(false) { - unload_library(library); - } -}; - -DG_DECLARE_STATIC_VAR_IN_CLASS(KernelRuntime, cuda_home); - -template -class LaunchRuntime { -public: - template - static std::string generate(const Args& args) { - auto code = Derived::generate_impl(args); - - // NOTES: we require that `generate_impl`'s includes never change - static std::string include_hash; - if (include_hash.empty()) - include_hash = include_parser->get_hash_value(code); - - // TODO: optimize string concat performance - code = fmt::format("// Includes' hash value: {}\n{}", include_hash, code); - if (get_env("DG_JIT_DEBUG")) - printf("Generated kernel code:\n%s\n", code.c_str()); - return code; - } - - template - static void launch(const std::shared_ptr& kernel_runtime, const Args& args) { - const auto kernel = kernel_runtime->kernel; - const auto stream = at::cuda::getCurrentCUDAStream(); - LaunchArgs launch_args = args.launch_args; - - // Allow runtime override from Python. - // NOTES: the default is enabled. - launch_args.enable_pdl = device_runtime->get_pdl(); - - const dim3 grid_dim = {static_cast(launch_args.grid_dim.first), - static_cast(launch_args.grid_dim.second), - 1}; - const dim3 block_dim = {static_cast(launch_args.num_threads), 1, 1}; - auto config = construct_launch_config(kernel, stream, launch_args.smem_size, - grid_dim, block_dim, launch_args.cluster_dim, launch_args.enable_pdl); - - // Launch in the derived class - if (get_env("DG_JIT_DEBUG")) { - printf("Launch kernel with {%d, %d} x %d, shared memory: %d bytes, cluster: %d, pdl: %d, stream: %ld\n", - launch_args.grid_dim.first, launch_args.grid_dim.second, launch_args.num_threads, - launch_args.smem_size, launch_args.cluster_dim, launch_args.enable_pdl, stream.id()); - } - Derived::launch_impl(kernel, config, args); - } -}; - -} // namespace deep_gemm diff --git a/csrc/jit_kernels/heuristics/common.hpp b/csrc/jit_kernels/heuristics/common.hpp index 2b79a8b74a..ad9f6a35b3 100644 --- a/csrc/jit_kernels/heuristics/common.hpp +++ b/csrc/jit_kernels/heuristics/common.hpp @@ -1,12 +1,16 @@ #pragma once +#include +#include +#include #include + #include +#include #include "config.hpp" #include "runtime.hpp" #include "../../utils/layout.hpp" -#include "../../utils/system.hpp" namespace deep_gemm { @@ -37,7 +41,7 @@ static GemmConfig get_best_config(const GemmDesc& desc) { }; // Print configs for the first time - if (get_env("DG_JIT_DEBUG") or get_env("DG_PRINT_CONFIGS")) { + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { std::stringstream ss; ss << desc; const auto key = ss.str(); diff --git a/csrc/jit_kernels/heuristics/config.hpp b/csrc/jit_kernels/heuristics/config.hpp index fe40777ad8..a2dc769665 100644 --- a/csrc/jit_kernels/heuristics/config.hpp +++ b/csrc/jit_kernels/heuristics/config.hpp @@ -32,8 +32,18 @@ struct GemmDesc { int get_expected_k() const { return expected_k > 0 ? expected_k : k; } int get_expected_num_groups() const { return expected_num_groups > 0 ? expected_num_groups : num_groups; } + bool is_mxf4_mma() const { + return a_dtype == kPackedFP4 and b_dtype == kPackedFP4; + } + MmaKind get_mma_kind() const { - return a_dtype == torch::kBFloat16 ? MmaKind::BF16 : MmaKind::MXFP8FP4; + if (a_dtype == torch::kBFloat16) + return MmaKind::BF16; + return is_mxf4_mma() ? MmaKind::MXF4 : MmaKind::MXFP8FP4; + } + + int get_smem_pack_factor() const { + return is_mxf4_mma() ? 2 : 1; } void check_validity() const { @@ -43,7 +53,10 @@ struct GemmDesc { DG_HOST_ASSERT(a_dtype == torch::kFloat8_e4m3fn or a_dtype == kPackedFP4); DG_HOST_ASSERT(b_dtype == torch::kFloat8_e4m3fn or b_dtype == kPackedFP4); } - DG_HOST_ASSERT(cd_dtype == torch::kBFloat16 or cd_dtype == torch::kFloat); + // FP8 D implies casting with dynamic per-32 UE8M0 SFD output, only exposed for batched GEMMs + DG_HOST_ASSERT(cd_dtype == torch::kBFloat16 or cd_dtype == torch::kFloat or + (cd_dtype == torch::kFloat8_e4m3fn and gemm_type == GemmType::Batched and + not with_accumulation)); DG_HOST_ASSERT(num_sms % 2 == 0); } @@ -110,11 +123,13 @@ struct StorageConfig { struct PipelineConfig { int smem_size; int num_stages; + int num_tma_store_stages; friend std::ostream& operator << (std::ostream& os, const PipelineConfig& config) { os << "PipelineConfig(" << "smem_size=" << config.smem_size - << ", num_stages=" << config.num_stages << ")"; + << ", num_stages=" << config.num_stages + << ", num_tma_store_stages=" << config.num_tma_store_stages << ")"; return os; } }; diff --git a/csrc/jit_kernels/heuristics/mega_gate.hpp b/csrc/jit_kernels/heuristics/mega_gate.hpp new file mode 100644 index 0000000000..b428c2e2b5 --- /dev/null +++ b/csrc/jit_kernels/heuristics/mega_gate.hpp @@ -0,0 +1,236 @@ +#pragma once + +#include +#include + +#include + +#include "../../utils/exception.hpp" +#include "../../utils/math.hpp" +#include "sm100.hpp" + +namespace deep_gemm { + +namespace mega_gate_layout = layout::mega_gate; + +constexpr int kNumMaxGateWarpgroups = 7; +constexpr int kNumTopkTokensPerWarpgroup = 4; +constexpr int kNumTransposeColsPerWarpgroup = 8; +constexpr int kMegaGateMaxSplitK = 8; +constexpr int kNumMaxStages = 32; +constexpr int kNumMaxRoutedExperts = 512; +constexpr int kNumStageBarrierBytes = 2 * static_cast(sizeof(uint64_t)); +constexpr int kNumFixedBarrierBytes = (2 * static_cast(mega_gate_layout::kNumEpilogueStages) + 1) * static_cast(sizeof(uint64_t)); + +static int get_num_fixed_smem_bytes(const int& num_routed_experts, const bool& has_bias, + const bool& has_image_token_mask, const bool& to_physical_map_exists) { + return align(kNumFixedBarrierBytes + static_cast(sizeof(uint32_t)) + + static_cast(mega_gate_layout::get_num_metadata_cache_bytes(num_routed_experts, has_bias, has_image_token_mask, to_physical_map_exists)), + static_cast(mega_gate_layout::kSharedMemoryAlignment)); +} + +static int get_num_stage_bytes(const int& load_block_m, const int& load_block_n) { + return (load_block_m + load_block_n) * static_cast(mega_gate_layout::BLOCK_K) * + static_cast(sizeof(cutlass::bfloat16_t)) + kNumStageBarrierBytes; +} + +static int get_num_stages(const int& smem_capacity, const int& num_fixed_smem_bytes, + const int& load_block_m, const int& load_block_n) { + return std::min(kNumMaxStages, + (smem_capacity - num_fixed_smem_bytes) / get_num_stage_bytes(load_block_m, load_block_n)); +} + +static int get_num_smem_bytes(const int& num_stages, const int& num_fixed_smem_bytes, + const int& load_block_m, const int& load_block_n) { + return align(num_fixed_smem_bytes + num_stages * get_num_stage_bytes(load_block_m, load_block_n), + static_cast(mega_gate_layout::kSharedMemoryAlignment)); +} + +struct SM100BF16MegaGateConfig { + int block_tokens; + int num_mma_ctas; + int num_split_k; + int num_expert_groups; + int num_stages; + int smem_size; + int num_gate_warpgroups; + int num_launch_sms; +}; + +static int get_num_worker_groups(const int& num_task_ctas, const int& num_sms) { + return num_sms / num_task_ctas; +} + +static int select_wave_tile(const int& num_tokens, const int& num_workers) { + const auto num_waves = ceil_div(ceil_div(num_tokens, + static_cast(mega_gate_layout::kNumMaxBlockTokens)), + num_workers); + return std::min(static_cast(mega_gate_layout::kNumMaxBlockTokens), + align(ceil_div(num_tokens, num_waves * num_workers), static_cast(mega_gate_layout::kNumMinBlockTokens))); +} + +static int get_num_token_cols(const int& block_tokens, const int& num_experts_per_group, const int& num_mma_ctas) { + return num_experts_per_group == static_cast(mega_gate_layout::kExpertAlignment) + ? block_tokens / num_mma_ctas : block_tokens; +} + +struct MegaGateConfig { + int num_expert_groups; + int num_mma_ctas; + int num_split_k; + int block_tokens; + int num_token_cols; + int num_launch_sms; + int num_waves; +}; + +static int get_num_waves(const int& num_tokens, const int& num_workers) { + return ceil_div(ceil_div(num_tokens, select_wave_tile(num_tokens, num_workers)), num_workers); +} + +static bool is_task_doubling_free(const int& num_tokens, const int& num_sms, const MegaGateConfig& candidate) { + const auto num_task_ctas = candidate.num_expert_groups * candidate.num_mma_ctas * candidate.num_split_k; + const auto num_doubled_task_ctas = num_task_ctas * 2; + if (num_doubled_task_ctas > num_sms or + num_doubled_task_ctas >= static_cast(mega_gate_layout::kNumMaxLogicalCtas)) + return false; + + const auto num_doubled_workers = get_num_worker_groups(num_doubled_task_ctas, num_sms); + if (candidate.num_mma_ctas == 1 and candidate.num_split_k != kMegaGateMaxSplitK) + return get_num_waves(num_tokens, num_doubled_workers) <= candidate.num_waves + 2; + if (candidate.num_mma_ctas == 2 and candidate.num_split_k == 1 and candidate.num_waves == 1) + return select_wave_tile(num_tokens, num_doubled_workers) == candidate.block_tokens; + return false; +} + +static std::vector get_mega_gate_candidates(const int& num_tokens, const int& hidden, + const int& num_routed_experts, const int& num_sms) { + constexpr int kExpertAlignment = mega_gate_layout::kExpertAlignment; + std::vector candidates; + for (int num_expert_groups = 1; num_expert_groups <= num_routed_experts / kExpertAlignment; ++ num_expert_groups) { + if (num_routed_experts % (kExpertAlignment * num_expert_groups) != 0) + continue; + + const auto num_experts_per_group = num_routed_experts / num_expert_groups; + if (num_experts_per_group > 2 * kExpertAlignment) + continue; + + for (int num_mma_ctas = 1; num_mma_ctas <= 2; ++ num_mma_ctas) { + if (num_experts_per_group > kExpertAlignment and num_mma_ctas != 2) + continue; + + for (int num_split_k = 1; num_split_k <= kMegaGateMaxSplitK; num_split_k *= 2) { + if (hidden % (num_split_k * static_cast(mega_gate_layout::BLOCK_K)) != 0) + continue; + + if (num_mma_ctas == 2 and num_split_k == kMegaGateMaxSplitK) + continue; + + const auto num_task_ctas = num_expert_groups * num_mma_ctas * num_split_k; + if (num_task_ctas > num_sms) + continue; + + if (num_task_ctas >= static_cast(mega_gate_layout::kNumMaxLogicalCtas)) + continue; + + const auto num_workers = get_num_worker_groups(num_task_ctas, num_sms); + const auto block_tokens = select_wave_tile(num_tokens, num_workers); + const auto num_token_blocks = ceil_div(num_tokens, block_tokens); + const auto num_token_cols = std::min(get_num_token_cols(block_tokens, num_experts_per_group, num_mma_ctas), + kNumMaxGateWarpgroups * kNumTransposeColsPerWarpgroup); + const MegaGateConfig candidate = {num_expert_groups, num_mma_ctas, num_split_k, + block_tokens, num_token_cols, + std::min(num_workers, num_token_blocks) * num_task_ctas, + ceil_div(num_token_blocks, num_workers)}; + if (is_task_doubling_free(num_tokens, num_sms, candidate)) + continue; + + candidates.push_back(candidate); + } + } + } + + DG_HOST_ASSERT(not candidates.empty()); + return candidates; +} + +static bool compare_mega_gate(const MegaGateConfig& a, const MegaGateConfig& b) { + if (a.num_token_cols != b.num_token_cols) + return a.num_token_cols > b.num_token_cols; + + if (a.num_launch_sms != b.num_launch_sms) + return a.num_launch_sms > b.num_launch_sms; + + if (a.num_mma_ctas != b.num_mma_ctas) + return a.num_mma_ctas > b.num_mma_ctas; + + if (a.num_split_k != b.num_split_k) + return a.num_split_k < b.num_split_k; + + return a.num_waves < b.num_waves; +} + +static MegaGateConfig select_mega_gate_config(const int& num_tokens, const int& hidden, + const int& num_routed_experts, const int& num_sms) { + const auto candidates = get_mega_gate_candidates(num_tokens, hidden, num_routed_experts, num_sms); + return *std::min_element(candidates.begin(), candidates.end(), compare_mega_gate); +} + +static int get_num_balanced_warpgroups(const int& num_units) { + return ceil_div(num_units, ceil_div(num_units, kNumMaxGateWarpgroups)); +} + +static int select_gate_warpgroups(const MegaGateConfig& config, const int& num_experts_per_group) { + if (config.num_waves > 1) + return kNumMaxGateWarpgroups; + + const auto num_task_ctas = config.num_expert_groups * config.num_mma_ctas * config.num_split_k; + const auto num_transpose_units = ceil_div(get_num_token_cols(config.block_tokens, num_experts_per_group, config.num_mma_ctas), + kNumTransposeColsPerWarpgroup); + const auto num_topk_units = ceil_div(config.block_tokens, kNumTopkTokensPerWarpgroup * num_task_ctas); + return std::max(get_num_balanced_warpgroups(num_transpose_units), + get_num_balanced_warpgroups(num_topk_units)); +} + +static SM100BF16MegaGateConfig get_sm100_bf16_mega_gate_config(const int& num_tokens, const int& hidden, + const int& num_actual_routed_experts, + const int& num_sms, const bool& has_bias, + const bool& has_image_token_mask, + const bool& to_physical_map_exists) { + DG_HOST_ASSERT(num_tokens > 0 and num_tokens <= static_cast(mega_gate_layout::kNumMaxTokens) and num_sms > 0); + DG_HOST_ASSERT(hidden > 0 and hidden % 256 == 0); + DG_HOST_ASSERT(num_actual_routed_experts > 0 and num_actual_routed_experts <= kNumMaxRoutedExperts and + num_actual_routed_experts % 4 == 0); + const auto num_routed_experts = align(num_actual_routed_experts, static_cast(mega_gate_layout::kExpertAlignment)); + auto config = select_mega_gate_config(num_tokens, hidden, num_routed_experts, num_sms); + if (heuristics_runtime->get_deterministic_algorithms()) + config.num_split_k = 1; + const auto token_mma = config.block_tokens; + + const auto experts_per_group = num_routed_experts / config.num_expert_groups; + const auto expert_mma = std::min(experts_per_group, 2 * static_cast(mega_gate_layout::kExpertAlignment)); + DG_HOST_ASSERT(expert_mma == static_cast(mega_gate_layout::kExpertAlignment) or config.num_mma_ctas == 2); + const auto load_block_m = token_mma / config.num_mma_ctas; + const auto load_block_n = expert_mma / config.num_mma_ctas; + const auto num_gate_warpgroups = select_gate_warpgroups(config, experts_per_group); + const auto num_fixed_smem_bytes = get_num_fixed_smem_bytes(num_routed_experts, has_bias, + has_image_token_mask, to_physical_map_exists); + const auto num_stages = get_num_stages(SM100ArchSpec::smem_capacity, num_fixed_smem_bytes, + load_block_m, load_block_n); + const auto smem_size = get_num_smem_bytes(num_stages, num_fixed_smem_bytes, + load_block_m, load_block_n); + DG_HOST_ASSERT(num_stages >= 2 and smem_size <= SM100ArchSpec::smem_capacity); + + return SM100BF16MegaGateConfig{ + .block_tokens = token_mma, + .num_mma_ctas = config.num_mma_ctas, + .num_split_k = config.num_split_k, + .num_expert_groups = config.num_expert_groups, + .num_stages = num_stages, + .smem_size = smem_size, + .num_gate_warpgroups = num_gate_warpgroups, + .num_launch_sms = config.num_launch_sms, + }; +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/heuristics/mega_moe.hpp b/csrc/jit_kernels/heuristics/mega_moe.hpp index 0a7a75794d..e018a33acf 100644 --- a/csrc/jit_kernels/heuristics/mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/mega_moe.hpp @@ -2,15 +2,19 @@ #include #include +#include +#include #include +#include #include +#include #include #include +#include #include "../../utils/exception.hpp" #include "../../utils/math.hpp" -#include "../../utils/system.hpp" #include "sm100.hpp" namespace deep_gemm { @@ -61,7 +65,7 @@ struct MegaMoEConfig { static MmaKind parse_mma_kind(const std::string& mma_type_str) { if (mma_type_str == "bf16xbf16") return MmaKind::BF16; - DG_HOST_ASSERT(mma_type_str == "fp8xfp4"); + DG_HOST_ASSERT(mma_type_str == "fp8xfp4" or mma_type_str == "fp8xfp8"); return MmaKind::MXFP8FP4; } @@ -78,29 +82,26 @@ static std::tuple get_block_config_for_mega_moe( const int& num_max_tokens_per_rank, const int& num_topk, const int& num_tokens, const MmaKind& mma_kind) { - auto [cluster_size, block_m, store_block_m, block_k, 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 and larger BLOCK_K for less synchronization - return {2, 16, 8, 256, 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, 128, 2}; - } else if (num_expected_tokens_per_expert <= 32.5) { - // Medium batch size, small EP, decoding, e.g. 6/384 experts, EP8, bsz 256 - return {2, 64, 32, 128, 1}; - } else if (num_expected_tokens_per_expert <= 64.5) { - // Large batch size, small EP, decoding, e.g. 6/384 experts, EP8, bsz 512 - return {2, 96, 16, 128, 2}; - } else if (num_expected_tokens_per_expert <= 96.5) { - // Medium batch size, Medium EP, decoding, e.g. 6/384 experts, EP16, bsz 256, or EP32, bsz128 - return {2, 128, 32, 128, 2}; - } else { - // Prefill, or large EP decoding - return {2, 192, 32, 128, 2}; + // Expected tokens per expert, plus a one-sigma routing margin for the tile choice + const float num_expected_tokens = static_cast(num_tokens) * num_ranks * num_topk / num_experts; + const float num_covered_tokens = num_expected_tokens + std::sqrt(num_expected_tokens); + + // Every M block costs roughly 128 extra rows (its tasks reload the weight tiles), so use the fewest blocks + // per expert and the smallest tile covering them: one block up to 192 rows, then blocks of up to 240 rows. + // A single 240-row block per expert only pays off with many local experts (about 14 or more blocks per rank). + int block_m = num_expected_tokens <= 10 ? 16 : 32; + if (num_expected_tokens > 24) { + int num_blocks = static_cast(std::ceil(num_covered_tokens / 240)); + if (num_blocks == 1 and num_covered_tokens > 192 and num_experts / num_ranks < 14) + num_blocks = 2; + for (const int& candidate: {64, 128, 192, 240}) { + block_m = candidate; + if (num_blocks * candidate >= num_covered_tokens) + break; } - }(); - block_k /= get_num_mma_elem_bytes(mma_kind); + } + const int store_block_m = block_m <= 16 ? 8 : block_m <= 64 ? 16 : block_m <= 192 ? 32 : 40; + const int block_k = 128 / get_num_mma_elem_bytes(mma_kind); // Check whether our `block_m` lies in `kCandidateBlockM` DG_HOST_ASSERT(std::any_of( @@ -108,8 +109,8 @@ static std::tuple get_block_config_for_mega_moe( [=](const auto& candidate) { return candidate == block_m; }) ); - // Return configs - return {cluster_size, block_m, store_block_m, block_k, num_epilogue_warpgroups * 128}; + // Return configs: 2-CTA clusters and 2 epilogue warpgroups + return {2, block_m, store_block_m, block_k, 2 * 128}; } static std::pair get_pipeline_config_for_mega_moe( @@ -191,12 +192,13 @@ static MegaMoEConfig get_mega_moe_config( // Block config const auto [cluster_size, block_m, store_block_m, block_k, num_epilogue_threads] = get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens, mma_kind); + DG_HOST_ASSERT(num_ring_tokens % block_m == 0); const int block_n = 128; const int load_block_m = block_m / 2; const int load_block_n = block_n; const auto [sf_block_m, sf_block_n] = is_mma_with_sf(mma_kind) ? SM100ArchSpec::get_sf_uttcp_aligned_block_sizes(block_m, block_n, MmaKind::MXFP8FP4) : std::pair(0, 0); - // NOTES: FP8 activations and FP4 weights (unpacked to 8-bit in smem) both use 128B swizzle + // NOTES: FP8 activations, FP8 weights, and FP4 weights unpacked to 8-bit in smem all use 128B swizzle const int swizzle_acts_mode = 128; const int swizzle_weights_mode = 128; const int gran_k = 32; @@ -205,10 +207,10 @@ static MegaMoEConfig get_mega_moe_config( const int num_dispatch_threads = 128; const int num_non_epilogue_threads = 128; - // Pull: divide token bytes by 2 until <= kPullThreshold - constexpr int kPullThreshold = 4096; + // Pull: divide token bytes by 2 until <= num_max_pull_bytes + const int num_max_pull_bytes = is_mma_with_sf(mma_kind) ? 8192 : 4096; int num_bytes_per_pull = hidden * get_num_mma_elem_bytes(mma_kind); - while (num_bytes_per_pull > kPullThreshold) { + while (num_bytes_per_pull > num_max_pull_bytes) { DG_HOST_ASSERT(num_bytes_per_pull % 2 == 0); num_bytes_per_pull /= 2; } @@ -234,8 +236,8 @@ static MegaMoEConfig get_mega_moe_config( }; // Print configs for the first time - if (get_env("DG_JIT_DEBUG") or get_env("DG_PRINT_CONFIGS")) { - const auto key = fmt::format( + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { + const auto key = std::format( "MegaMoEConfig(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; diff --git a/csrc/jit_kernels/heuristics/runtime.hpp b/csrc/jit_kernels/heuristics/runtime.hpp index 6f84648417..417c612421 100644 --- a/csrc/jit_kernels/heuristics/runtime.hpp +++ b/csrc/jit_kernels/heuristics/runtime.hpp @@ -1,20 +1,30 @@ #pragma once -#include "../../jit/device_runtime.hpp" +#include + +#include "../../runtime/jit.hpp" #include "../../utils/exception.hpp" -#include "../../utils/lazy_init.hpp" namespace deep_gemm { class HeuristicsRuntime { +public: static constexpr int kLegacyMKAlignmentForContiguousLayout = 128; bool ignore_compile_dims = false; + bool deterministic_algorithms = false; int block_m_multiple_of = 1; int block_n_multiple_of = 1; int mk_alignment_for_contiguous_layout = kLegacyMKAlignmentForContiguousLayout; -public: + void use_deterministic_algorithms(const bool enabled) { + deterministic_algorithms = enabled; + } + + bool get_deterministic_algorithms() const { + return deterministic_algorithms; + } + void set_ignore_compile_dims(const bool& new_value) { ignore_compile_dims = new_value; } @@ -45,18 +55,15 @@ class HeuristicsRuntime { } static int get_theoretical_mk_alignment_for_contiguous_layout(const std::optional& expected_m) { - if (device_runtime->get_arch_major() != 10) + if (jit->device.get_arch_major() != 10) return kLegacyMKAlignmentForContiguousLayout; - int block_m = 224, mma_step = 32; - if (expected_m.has_value()) { - // Reduce `block_m` while ensuring it covers `m` - for (; block_m > 32 and block_m - mma_step >= expected_m.value(); block_m -= mma_step); - } - return block_m; + // The newly supported UMMA_N=256 allows a fixed alignment of 256, which is friendlier to MoE cast, M-grouped, and K-grouped operators. + // NOTES: `expected_m` is ignored, so small values may incur performance loss; use Mega MoE directly for such workloads. + return 256; } }; -static auto heuristics_runtime = LazyInit([](){ return std::make_shared(); }); +inline auto heuristics_runtime = deep_jit::LazyInit([](){ return std::make_shared(); }); } // namespace deep_gemm diff --git a/csrc/jit_kernels/heuristics/sm100.hpp b/csrc/jit_kernels/heuristics/sm100.hpp index c8e9e2e07f..3011c18ab0 100644 --- a/csrc/jit_kernels/heuristics/sm100.hpp +++ b/csrc/jit_kernels/heuristics/sm100.hpp @@ -1,5 +1,6 @@ #pragma once +#include #include // Reuse some types in the JIT modules #include @@ -18,15 +19,17 @@ struct SM100ArchSpec { const int& block_m, const int& block_n, const MmaKind& mma_kind) { constexpr int num_utccp_aligned_elems = 128; switch (mma_kind) { - case MmaKind::BF16: return {0, 0}; - case MmaKind::MXFP8FP4: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)}; + case MmaKind::BF16: return {0, 0}; + case MmaKind::MXFP8FP4: + case MmaKind::MXF4: + return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)}; default: DG_HOST_UNREACHABLE("Unknown dtype"); } } static std::vector get_layout_candidates(const GemmDesc& desc) { // Block K is always in a fixed manner - const int block_k = 128 / get_element_size(desc.get_mma_kind()); + const int block_k = 128 * 8 / get_num_element_bits(desc.get_mma_kind()); // Always enable swap A/B (and multicasting if possible) for m-grouped GEMMs if (desc.gemm_type == GemmType::MGroupedContiguous or @@ -122,9 +125,10 @@ struct SM100ArchSpec { // Check tensor memory capacity const auto [sf_block_m, sf_block_n] = get_sf_uttcp_aligned_block_sizes(block_m, block_n, desc.get_mma_kind()); - const auto tmem_sf_cols = desc.get_mma_kind() == MmaKind::MXFP8FP4 ? sf_block_m / 32 + sf_block_n / 32 : 0; + const auto sf_block_k = block_k / 128; + const auto tmem_sf_cols = desc.get_mma_kind() == MmaKind::BF16 ? 0 : sf_block_m * sf_block_k / 32 + sf_block_n * sf_block_k / 32; const auto umma_n = swap_ab ? block_m : block_n; - if (2 * umma_n + tmem_sf_cols > 512) + if (umma_n + tmem_sf_cols > 512) continue; const auto layout = Layout{swap_ab, block_m, block_n, block_k, cluster_m, cluster_n}; @@ -143,6 +147,17 @@ struct SM100ArchSpec { } } + // Dynamic FP8 output requires the physical store atom to own complete per-32 SF groups + if (desc.cd_dtype == torch::kFloat8_e4m3fn) { + const auto owns_partial_sf_groups = [&](const Layout& layout) { + const auto store_block_n = layout.swap_ab ? layout.block_n : + get_storage_config(desc, layout).swizzle_cd_mode; + return store_block_n % 32 != 0; + }; + candidates.erase(std::remove_if(candidates.begin(), candidates.end(), owns_partial_sf_groups), + candidates.end()); + } + DG_HOST_ASSERT(not candidates.empty()); return candidates; } @@ -158,11 +173,13 @@ struct SM100ArchSpec { const auto store_block_n = layout.block_n; // Decide swizzling by the inner dim - // TODO: support FP4 sub-byte - const auto swizzle_mode_a = get_swizzle_mode( - desc.major_a == cute::UMMA::Major::K ? layout.block_k : load_block_m, c10::elementSize(desc.a_dtype)); - const auto swizzle_mode_b = get_swizzle_mode( - desc.major_b == cute::UMMA::Major::K ? layout.block_k : load_block_n, c10::elementSize(desc.b_dtype)); + const int pack_factor = desc.get_smem_pack_factor(); + const int swizzle_mode_a = get_swizzle_mode( + (desc.major_a == cute::UMMA::Major::K ? layout.block_k : load_block_m) / pack_factor, + c10::elementSize(desc.a_dtype)); + const int swizzle_mode_b = get_swizzle_mode( + (desc.major_b == cute::UMMA::Major::K ? layout.block_k : load_block_n) / pack_factor, + c10::elementSize(desc.b_dtype)); const auto swizzle_mode_cd = get_swizzle_mode( store_block_n, c10::elementSize(desc.cd_dtype)); @@ -173,26 +190,43 @@ struct SM100ArchSpec { }; } + static int get_num_tma_store_stages(const GemmDesc& desc, const Layout& layout) { + // With single-stage TMA stores, each SM keeps at most one bulk store in + // flight. This paces the store traffic and improves the achieved DRAM + // throughput when C/D flushes dominate DRAM traffic; it also frees up + // shared memory for the A/B mainloop and halves the C/D smem read/write + // footprint. Only enabled for k-grouped GEMMs. + if (desc.get_mma_kind() == MmaKind::BF16 or layout.swap_ab or + not is_k_grouped_contiguous(desc.gemm_type)) + return 2; + + const int num_k_blocks_per_group = desc.k / std::max(desc.num_groups, 1) / layout.block_k; + const int min_k_blocks = desc.cd_dtype != torch::kFloat ? 16 : + desc.with_accumulation ? 24 : 32; + return num_k_blocks_per_group >= min_k_blocks ? 1 : 2; + } + static PipelineConfig get_pipeline_config(const GemmDesc& desc, const Layout& layout, const StorageConfig& storage_config) { constexpr int kNumMaxStages = 32; // C/D for TMA stores - const int smem_cd = layout.swap_ab ? storage_config.store_block_m * storage_config.store_block_n * c10::elementSize(desc.cd_dtype) * 2 - : storage_config.store_block_m * storage_config.swizzle_cd_mode * 2; + const int num_tma_store_stages = get_num_tma_store_stages(desc, layout); + const int smem_cd = (layout.swap_ab ? storage_config.store_block_m * storage_config.store_block_n * c10::elementSize(desc.cd_dtype) + : storage_config.store_block_m * storage_config.swizzle_cd_mode) * num_tma_store_stages; // TODO: remove SF barriers for BF16 GEMMs - // TMA full/empty barriers, with-SF full barriers, tensor memory full/empty barriers + // A/B-and-SF-transpose full, SF-TMA full, empty, and tensor memory full/empty/overlap barriers // NOTES: some shapes may only have 1 epilogue stage, but we still allocate space for 2 stages // NOTES: the last barrier is for tensor core utilization control - const int smem_barriers = kNumMaxStages * 8 * 3 + 2 * 8 * 2 + 8; + const int smem_barriers = kNumMaxStages * 8 * 3 + 2 * 8 * 3 + 8; // Tensor memory pointer const int smem_tmem_ptr = 4; // Calculate A/B per stages - // TODO: consider FP4 - const int smem_a_per_stage = storage_config.load_block_m * layout.block_k * c10::elementSize(desc.a_dtype); - const int smem_b_per_stage = storage_config.load_block_n * layout.block_k * c10::elementSize(desc.b_dtype); + const int pack_factor = desc.get_smem_pack_factor(); + const int smem_a_per_stage = storage_config.load_block_m * layout.block_k * c10::elementSize(desc.a_dtype) / pack_factor; + const int smem_b_per_stage = storage_config.load_block_n * layout.block_k * c10::elementSize(desc.b_dtype) / pack_factor; // Calculate SF A/B per stages int smem_sfa_per_stage = 0; @@ -200,8 +234,8 @@ struct SM100ArchSpec { if (desc.kernel_type == KernelType::Kernel1D1D) { const auto [sf_block_m, sf_block_n] = get_sf_uttcp_aligned_block_sizes( layout.block_m, layout.block_n, desc.get_mma_kind()); - smem_sfa_per_stage = sf_block_m * 4; - smem_sfb_per_stage = sf_block_n * 4; + smem_sfa_per_stage = sf_block_m * layout.block_k / 32; + smem_sfb_per_stage = sf_block_n * layout.block_k / 32; } // Calculate stages @@ -212,7 +246,8 @@ struct SM100ArchSpec { kNumMaxStages); return { smem_extra + num_stages * smem_per_stage, - num_stages + num_stages, + num_tma_store_stages }; } diff --git a/csrc/jit_kernels/heuristics/sm90.hpp b/csrc/jit_kernels/heuristics/sm90.hpp index c411fb7e01..144bfd19e3 100644 --- a/csrc/jit_kernels/heuristics/sm90.hpp +++ b/csrc/jit_kernels/heuristics/sm90.hpp @@ -57,7 +57,7 @@ struct SM90ArchSpec { block_n_candidates.push_back(i); // Block K is always in a fixed manner - const int block_k = 128 / get_element_size(desc.get_mma_kind()); + const int block_k = 128 * 8 / get_num_element_bits(desc.get_mma_kind()); // Disable multicast for performance const bool disable_multicast = @@ -182,7 +182,8 @@ struct SM90ArchSpec { kNumMaxStages); return { smem_extra + num_stages * smem_per_stage, - num_stages + num_stages, + 2 }; } diff --git a/csrc/jit_kernels/heuristics/utils.hpp b/csrc/jit_kernels/heuristics/utils.hpp index 17d2ae07e6..fc643880c1 100644 --- a/csrc/jit_kernels/heuristics/utils.hpp +++ b/csrc/jit_kernels/heuristics/utils.hpp @@ -9,6 +9,15 @@ namespace deep_gemm { +static int get_num_element_bits(const MmaKind& mma_kind) { + switch (mma_kind) { + case MmaKind::BF16: return 16; + case MmaKind::MXFP8FP4: return 8; + case MmaKind::MXF4: return 4; + default: DG_HOST_UNREACHABLE("Unknown MMA kind"); + } +} + template static int get_swizzle_mode(const int& block_size, const size_type_t& elem_size) { // `> 0` means interleaving diff --git a/csrc/jit_kernels/impls/epilogue.hpp b/csrc/jit_kernels/impls/epilogue.hpp index 1003df4c3f..c132a28942 100644 --- a/csrc/jit_kernels/impls/epilogue.hpp +++ b/csrc/jit_kernels/impls/epilogue.hpp @@ -3,10 +3,43 @@ #include #include +#include + +#include + +#include "../../utils/exception.hpp" + namespace deep_gemm { static std::string get_default_epilogue_type(const std::optional& epilogue_type) { return epilogue_type.value_or("epilogue::transform::EpilogueIdentity"); } +// The full host-determined epilogue input: the operator type to instantiate the kernel with, +// and the operator's runtime state to launch the kernel with +struct EpilogueInput { + std::string type = "epilogue::transform::EpilogueIdentity"; + EpilogueArgs args; +}; + +// Construct the epilogue input, keeping the operator type and its state consistent +// NOTES: the selection only depends on which inputs are given, never on their values +static EpilogueInput make_epilogue_input(const int& m, const int& n, + const std::optional& epilogue_type = std::nullopt, + const std::optional& alpha = std::nullopt, + const std::optional& sfd = std::nullopt) { + if (sfd.has_value()) { + DG_HOST_ASSERT(not epilogue_type.has_value() and not alpha.has_value()); + return {"epilogue::transform::EpilogueDynamicScaledFP8", + {.sfd = static_cast(sfd->data_ptr()), + .sfd_stride = static_cast(sfd->stride(-1)), + .shape_m = static_cast(m), .shape_n = static_cast(n)}}; + } + if (alpha.has_value()) { + DG_HOST_ASSERT(not epilogue_type.has_value()); + return {"epilogue::transform::EpilogueWithAlpha", {.alpha = alpha.value()}}; + } + return {get_default_epilogue_type(epilogue_type), {}}; +} + } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/runtime_utils.hpp b/csrc/jit_kernels/impls/runtime_utils.hpp index 2e617d11a5..15cd7a14f7 100644 --- a/csrc/jit_kernels/impls/runtime_utils.hpp +++ b/csrc/jit_kernels/impls/runtime_utils.hpp @@ -1,12 +1,20 @@ #pragma once +#include +#include +#include +#include +#include +#include + #include #include +#include +#include + #include "../heuristics/sm90.hpp" -#include "../../jit/handle.hpp" #include "../../utils/math.hpp" -#include "../../utils/system.hpp" #include "../../utils/exception.hpp" namespace deep_gemm { @@ -51,20 +59,21 @@ static std::string to_string(const GemmType& type) { DG_HOST_UNREACHABLE("Unknown GEMM type"); } -static std::string to_string(const at::ScalarType& dtype) { +static std::string to_string(const at::ScalarType& dtype, const bool& fp4_unpacked_smem = true) { switch (dtype) { case torch::kInt: return "int"; case torch::kFloat: return "float"; case torch::kBFloat16: return "cutlass::bfloat16_t"; case torch::kFloat8_e4m3fn: return "cutlass::float_e4m3_t"; - case kPackedFP4: return "cutlass::detail::float_e2m1_unpacksmem_t"; + case kPackedFP4: return fp4_unpacked_smem ? "cutlass::detail::float_e2m1_unpacksmem_t" + : "cutlass::float_e2m1_t"; default: DG_HOST_UNREACHABLE("Unsupported dtype"); } } static std::string to_string(const float& v) { if (std::isfinite(v)) { - return fmt::format(R"({:a}f)", v); + return std::format("{}0x{:a}f", std::signbit(v) ? "-" : "", std::abs(v)); } else if (std::isinf(v)) { return v > 0 ? "cute::numeric_limits::infinity()" : "-cute::numeric_limits::infinity()"; @@ -83,21 +92,17 @@ static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType& case torch::kFloat: return CU_TENSOR_MAP_DATA_TYPE_FLOAT32; case torch::kBFloat16: return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; case torch::kFloat8_e4m3fn: return CU_TENSOR_MAP_DATA_TYPE_UINT8; -#if CUDA_VERSION >= 12080 case kPackedFP4: return fp4_unpacked_smem ? CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B : CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; -#endif default: DG_HOST_UNREACHABLE("Unsupported dtype"); } } static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const int& base) { -#if CUDA_VERSION >= 12080 if (base != 0) { DG_HOST_ASSERT(base == 32 and mode == 128); return CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B; } -#endif DG_HOST_ASSERT(base == 0); switch (mode) { @@ -118,6 +123,9 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t, const bool& allow_tf32 = false, const bool& fp4_unpacked_smem = true) { const auto elem_size = static_cast(t.element_size()); + const auto num_gmem_outer_stride_bytes = static_cast(gmem_outer_stride) * static_cast(elem_size); + DG_HOST_ASSERT(reinterpret_cast(t.data_ptr()) % 16u == 0u); + DG_HOST_ASSERT(num_gmem_outer_stride_bytes % 16u == 0u); if (swizzle_mode != 0) smem_inner_dim = swizzle_mode / elem_size; @@ -135,13 +143,13 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t, const cuuint32_t smem_dims[2] = {static_cast(smem_inner_dim), static_cast(smem_outer_dim)}; const cuuint64_t gmem_strides[1] = {static_cast(gmem_outer_stride * elem_size), }; const cuuint32_t elem_strides[2] = {1, 1}; - if (get_env("DG_JIT_DEBUG")) { + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { printf("Making TMA desc: global memory: %d %d, shared memory: %d %d, outer stride: %d, swizzle: %d (base: %d), elem size: %d, pointer: %llu\n", gmem_inner_dim, gmem_outer_dim, smem_inner_dim, smem_outer_dim, gmem_outer_stride, swizzle_mode, swizzle_base, elem_size, reinterpret_cast(t.data_ptr())); } - DG_CUDA_DRIVER_CHECK(lazy_cuTensorMapEncodeTiled( + DJ_CUDA_DRIVER_CHECK(deep_jit::cuda::driver::lazy_cuTensorMapEncodeTiled( &tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem), 2, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base), @@ -157,6 +165,11 @@ static CUtensorMap make_tma_3d_desc(const torch::Tensor& t, const bool& allow_tf32 = false, const bool& fp4_unpacked_smem = true) { const auto elem_size = static_cast(t.element_size()); + const auto num_gmem_stride_0_bytes = static_cast(gmem_stride_0) * static_cast(elem_size); + const auto num_gmem_stride_1_bytes = static_cast(gmem_stride_1) * static_cast(elem_size); + DG_HOST_ASSERT(reinterpret_cast(t.data_ptr()) % 16u == 0u); + DG_HOST_ASSERT(num_gmem_stride_0_bytes % 16u == 0u); + DG_HOST_ASSERT(num_gmem_stride_1_bytes % 16u == 0u); if (swizzle_mode != 0) smem_dim_0 = swizzle_mode / elem_size; @@ -174,12 +187,12 @@ static CUtensorMap make_tma_3d_desc(const torch::Tensor& t, const cuuint32_t smem_dims[3] = {static_cast(smem_dim_0), static_cast(smem_dim_1), static_cast(smem_dim_2)}; const cuuint64_t gmem_strides[2] = {static_cast(gmem_stride_0 * elem_size), static_cast(gmem_stride_1 * elem_size)}; const cuuint32_t elem_strides[3] = {1, 1, 1}; - if (get_env("DG_JIT_DEBUG")) { + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { printf("Making 3D TMA desc: global memory: %d %d %d, shared memory: %d %d %d, outer stride: %d %d, swizzle: %d, elem size: %d\n", gmem_dim_0, gmem_dim_1, gmem_dim_2, smem_dim_0, smem_dim_1, smem_dim_2, gmem_stride_0, gmem_stride_1, swizzle_mode, elem_size); } - DG_CUDA_DRIVER_CHECK(lazy_cuTensorMapEncodeTiled( + DJ_CUDA_DRIVER_CHECK(deep_jit::cuda::driver::lazy_cuTensorMapEncodeTiled( &tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem), 3, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides, CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base), @@ -194,7 +207,8 @@ static CUtensorMap make_tma_a_desc(const cute::UMMA::Major& major, const int& outer_stride, const int& num_groups, const int& swizzle_mode, const int& swizzle_base = 0, - const bool& allow_tf32 = false) { + const bool& allow_tf32 = false, + const bool& fp4_unpacked_smem = true) { if (num_groups > 1) DG_HOST_ASSERT(major == cute::UMMA::Major::K); const auto [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_m * num_groups); @@ -204,7 +218,7 @@ static CUtensorMap make_tma_a_desc(const cute::UMMA::Major& major, smem_inner_dim, smem_outer_dim, outer_stride, swizzle_mode, swizzle_base, - allow_tf32); + allow_tf32, fp4_unpacked_smem); } static CUtensorMap make_tma_b_desc(const cute::UMMA::Major& major, @@ -214,7 +228,8 @@ static CUtensorMap make_tma_b_desc(const cute::UMMA::Major& major, const int& outer_stride, const int& num_groups, const int& swizzle_mode, const int& swizzle_base = 0, - const bool& allow_tf32 = false) { + const bool& allow_tf32 = false, + const bool& fp4_unpacked_smem = true) { const auto [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_n); const auto [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_n); @@ -224,7 +239,7 @@ static CUtensorMap make_tma_b_desc(const cute::UMMA::Major& major, smem_inner_dim, smem_outer_dim, outer_stride, swizzle_mode, swizzle_base, - allow_tf32); + allow_tf32, fp4_unpacked_smem); } static CUtensorMap make_tma_cd_desc(const torch::Tensor& t, diff --git a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp index f9b2f361cb..030e77c503 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_gemm.hpp @@ -1,33 +1,33 @@ #pragma once +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../../utils/math.hpp" #include "../heuristics/sm100.hpp" +#include "epilogue.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM100BF16GemmRuntime final: public LaunchRuntime { +class SM100BF16GemmRuntime final { public: struct Args { GemmDesc gemm_desc; GemmConfig gemm_config; - LaunchArgs launch_args; + deep_jit::cuda::LaunchOptions options; void* grouped_layout; CUtensorMap tensor_map_a; CUtensorMap tensor_map_b; CUtensorMap tensor_map_cd; + EpilogueInput epilogue; }; - static std::string generate_impl(const Args& args) { - return fmt::format(R"( + static void compile_and_launch(const std::string& tag, const Args& args) { + const auto kernel = jit->compile(tag, std::format(R"( #include using namespace deep_gemm; @@ -44,9 +44,9 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, - {}, - {}, + {}, {}, {}, {}, {}, + {}, {} >); }}; @@ -64,16 +64,18 @@ static void __instantiate_kernel() {{ args.gemm_config.launch_config.num_sms, heuristics_runtime->get_mk_alignment_for_contiguous_layout(), args.gemm_config.layout.swap_ab, args.gemm_desc.ensure_zero_padding, - to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation, to_string(args.gemm_desc.cd_dtype), - args.gemm_desc.tc_util); - } + to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation, + to_string(args.gemm_desc.cd_dtype), + args.epilogue.type, + args.gemm_desc.tc_util)); - 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.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, + // Launch + jit->launch( + kernel, args.options, + args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, args.epilogue.args, args.tensor_map_a, args.tensor_map_b, - args.tensor_map_cd)); + args.tensor_map_cd + ); } }; @@ -83,7 +85,8 @@ static void sm100_bf16_gemm(const torch::Tensor& a, const torch::Tensor& d, const int& m, const int& n, const int& k, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, - const std::string& compiled_dims) { + const std::string& compiled_dims, + const std::optional& alpha) { const auto desc = GemmDesc { .gemm_type = GemmType::Normal, .kernel_type = KernelType::KernelNoSF, @@ -92,8 +95,9 @@ static void sm100_bf16_gemm(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -113,21 +117,22 @@ static void sm100_bf16_gemm(const torch::Tensor& a, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); - // Launch - const SM100BF16GemmRuntime::Args& args = { + // Compile and launch + SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_gemm", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, - .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm100_bf16_gemm", code); - SM100BF16GemmRuntime::launch(runtime, args); + .tensor_map_cd = tensor_map_cd, + .epilogue = make_epilogue_input(m, n, std::nullopt, alpha) + }); } static void sm100_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, @@ -157,8 +162,9 @@ static void sm100_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, .ensure_zero_padding = ensure_zero_padding, .expected_m = expected_m_for_psum_layout.value_or(m), .expected_n = n, .expected_k = k, @@ -182,21 +188,21 @@ static void sm100_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); - // Launch - const SM100BF16GemmRuntime::Args args = { + // Compile and launch + SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_m_grouped_gemm_contiguous", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = grouped_layout.data_ptr(), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm100_bf16_m_grouped_gemm_contiguous", code); - SM100BF16GemmRuntime::launch(runtime, args); + }); } static void sm100_m_grouped_bf16_gemm_masked(const torch::Tensor& a, @@ -215,8 +221,9 @@ static void sm100_m_grouped_bf16_gemm_masked(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups }; const auto config = get_best_config(desc); @@ -237,21 +244,21 @@ static void sm100_m_grouped_bf16_gemm_masked(const torch::Tensor& a, static_cast(d.stride(-2)), num_groups, config.storage_config.swizzle_cd_mode); - // Launch - const SM100BF16GemmRuntime::Args args = { + // Compile and launch + SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_m_grouped_gemm_masked", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = masked_m.data_ptr(), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm100_bf16_m_grouped_gemm_masked", code); - SM100BF16GemmRuntime::launch(runtime, args); + }); } static void sm100_bf16_k_grouped_gemm(const torch::Tensor& a, @@ -276,8 +283,9 @@ static void sm100_bf16_k_grouped_gemm(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, // NOTES: expected_k is not used in SM100 get_best_config yet. .expected_m = m, .expected_n = n, .expected_k = expected_k, .expected_num_groups = num_groups }; @@ -294,27 +302,28 @@ static void sm100_bf16_k_grouped_gemm(const torch::Tensor& a, config.layout.block_k, static_cast(b.stride(0)), 1, config.storage_config.swizzle_b_mode); - const auto tensor_map_cd = make_tma_cd_desc(d, m, n, - config.storage_config.store_block_m, + const auto tensor_map_cd = make_tma_3d_desc(d, n, m, num_groups, config.storage_config.store_block_n, - static_cast(d.stride(1)), num_groups, + config.storage_config.store_block_m, 1, + static_cast(d.stride(1)), + static_cast(d.stride(0)), config.storage_config.swizzle_cd_mode); - // Launch kernel - const SM100BF16GemmRuntime::Args& args = { + // Compile and launch + SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_k_grouped_gemm", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = grouped_layout.data_ptr(), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm100_bf16_k_grouped_gemm", code); - SM100BF16GemmRuntime::launch(runtime, args); + }); } static void sm100_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, @@ -330,8 +339,9 @@ static void sm100_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, .cd_dtype = tensor_d.scalar_type(), .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -348,21 +358,21 @@ static void sm100_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, tensor_d.stride(0), tensor_d.stride(1), config.storage_config.swizzle_cd_mode); - // Launch - const SM100BF16GemmRuntime::Args& args = { + // Compile and launch + SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_bhr_hdr_bhd", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm100_bf16_bhr_hdr_bhd", code); - SM100BF16GemmRuntime::launch(runtime, args); + }); } static void sm100_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, @@ -378,8 +388,9 @@ static void sm100_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, .cd_dtype = tensor_d.scalar_type(), .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::MN, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -396,21 +407,21 @@ static void sm100_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, tensor_d.stride(0), tensor_d.stride(1), config.storage_config.swizzle_cd_mode); - // Launch - const SM100BF16GemmRuntime::Args& args = { + // Compile and launch + SM100BF16GemmRuntime::compile_and_launch("sm100_bf16_bhd_hdr_bhr", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm100_bf16_bhd_hdr_bhr", code); - SM100BF16GemmRuntime::launch(runtime, args); + }); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_gate.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_gate.hpp new file mode 100644 index 0000000000..e3c263ea2d --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_gate.hpp @@ -0,0 +1,123 @@ +#pragma once + +#include +#include +#include + +#include + +#include "../../runtime/jit.hpp" +#include "../../utils/exception.hpp" +#include "../heuristics/mega_gate.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +namespace mega_gate_layout = layout::mega_gate; + +static void sm100_bf16_mega_gate(const torch::Tensor& x, const torch::Tensor& weight, + const std::optional& bias, + const std::optional& image_bias, + const std::optional& image_token_mask, + const std::optional& mask, + const std::optional& fix_routing_mask, + const std::optional& to_physical_map, + const std::optional& logical_count, + const torch::Tensor& topk_idx, + const std::optional& unmapped_topk_idx, + const torch::Tensor& topk_weights, + const std::optional& force_random, + const int& num_tokens, const int& hidden, + const int& num_routed_experts, const int& num_topk, + const int& num_shared_experts, const float& routed_scaling_factor, + const int& ep_rank, const int& scoring_type, + const SM100BF16MegaGateConfig& config, + const torch::Tensor& scratch, + const torch::Tensor& score_barriers) { + const auto num_duplicate_experts = to_physical_map.has_value() ? static_cast(to_physical_map->size(1)) : 0; + const auto unmapped_topk_idx_stride = unmapped_topk_idx.has_value() ? unmapped_topk_idx->stride(0) : 0; + const auto num_aligned_experts = align(num_routed_experts, static_cast(mega_gate_layout::kExpertAlignment)); + const auto experts_per_group = num_aligned_experts / config.num_expert_groups; + const auto expert_mma = std::min(experts_per_group, 256); + const auto load_block_m = config.block_tokens / config.num_mma_ctas; + const auto load_block_n = expert_mma / config.num_mma_ctas; + const auto tensor_map_x = make_tma_a_desc(cute::UMMA::Major::K, x, num_tokens, hidden, + load_block_m, mega_gate_layout::BLOCK_K, + static_cast(x.stride(0)), 1, 128); + const auto tensor_map_weight = make_tma_b_desc(cute::UMMA::Major::K, weight, num_routed_experts, hidden, + load_block_n, mega_gate_layout::BLOCK_K, + static_cast(weight.stride(0)), 1, 128); + + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { + printf("Mega gate: T=%d, H=%d, E=%d, block_tokens=%d, mma_ctas=%d, split_k=%d, expert_groups=%d, " + "stages=%d, gate_warpgroups=%d, threads=%d, shared memory=%d\n", + num_tokens, hidden, num_routed_experts, config.block_tokens, config.num_mma_ctas, + config.num_split_k, config.num_expert_groups, config.num_stages, + config.num_gate_warpgroups, + mega_gate_layout::kNumNonEpilogueThreads + config.num_gate_warpgroups * 128, + config.smem_size); + } + + // Compile + const auto kernel = jit->compile("sm100_bf16_mega_gate", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_bf16_mega_gate_impl< + {}, {}, + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, {}, {}, {}, {}, {}, {}, {} + >); +}}; +)", + hidden, num_aligned_experts, + config.block_tokens, + config.num_stages, + config.num_gate_warpgroups * 128, + config.num_mma_ctas, config.num_split_k, config.num_expert_groups, + config.num_launch_sms, + num_topk, scoring_type, + num_routed_experts == num_aligned_experts, + (mask ? mask->data_ptr() : nullptr) != nullptr, + (unmapped_topk_idx ? unmapped_topk_idx->data_ptr() : nullptr) != nullptr, + (to_physical_map ? to_physical_map->data_ptr() : nullptr) != nullptr, + (bias ? bias->data_ptr() : nullptr) != nullptr, + (image_token_mask ? image_token_mask->data_ptr() : nullptr) != nullptr, + (fix_routing_mask ? fix_routing_mask->data_ptr() : nullptr) != nullptr, + (force_random ? force_random->data_ptr() : nullptr) != nullptr)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = config.smem_size, + .grid_dim = dim3(config.num_launch_sms, 1, 1), + .block_dim = dim3(mega_gate_layout::kNumNonEpilogueThreads + config.num_gate_warpgroups * 128, 1, 1), + .cluster_dim = dim3(config.num_mma_ctas, 1, 1), + }, + tensor_map_x, tensor_map_weight, + bias ? bias->data_ptr() : nullptr, + image_bias ? image_bias->data_ptr() : nullptr, + image_token_mask ? image_token_mask->data_ptr() : nullptr, + mask ? mask->data_ptr() : nullptr, + to_physical_map ? to_physical_map->data_ptr() : nullptr, + logical_count ? logical_count->data_ptr() : nullptr, + topk_idx.data_ptr(), + unmapped_topk_idx ? unmapped_topk_idx->data_ptr() : nullptr, + topk_weights.data_ptr(), + scratch.data_ptr(), score_barriers.data_ptr(), + static_cast(num_tokens), + static_cast(num_routed_experts), + static_cast(num_shared_experts), + static_cast(num_duplicate_experts), + routed_scaling_factor, static_cast(ep_rank), + unmapped_topk_idx_stride, + fix_routing_mask ? fix_routing_mask->data_ptr() : nullptr, + force_random ? force_random->data_ptr() : nullptr); +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp index 273874cfd9..6cf19427cd 100644 --- a/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_bf16_mega_moe.hpp @@ -1,113 +1,18 @@ #pragma once +#include #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 "../../runtime/runtime.hpp" +#include "../../utils/exception.hpp" #include "../heuristics/mega_moe.hpp" +#include "runtime_utils.hpp" namespace deep_gemm { -class SM100BF16MegaMoERuntime final : public LaunchRuntime { -public: - struct Args { - // Templated arguments - int num_max_tokens_per_rank; - int hidden, intermediate_hidden; - int num_experts, num_shared_experts, num_topk; - int num_ranks; - float activation_clamp; - bool fast_math; - MegaMoEConfig config; - - // Runtime arguments - void* y; - int* cumulative_local_expert_recv_stats; - int num_tokens; - layout::SymBuffer<> sym_buffer_ptrs; - - // Tensormap - CUtensorMap tensor_map_l1_acts; - CUtensorMap tensor_map_l1_weights; - CUtensorMap tensor_map_l1_output; - CUtensorMap tensor_map_l2_acts; - CUtensorMap tensor_map_l2_weights; - CUtensorMap tensor_map_shared_l1_acts; - CUtensorMap tensor_map_shared_l1_weights; - CUtensorMap tensor_map_shared_l1_output; - CUtensorMap tensor_map_shared_l2_acts; - CUtensorMap tensor_map_shared_l2_weights; - - // 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(&sm100_bf16_mega_moe_impl< - {}, - {}, {}, - {}, {}, - {}, {}, {}, - {}, - {}, - {}, - {}, - {}, - {}, {}, {}, - {}, {}, - {}, - {} - >); -}}; -)", args.num_max_tokens_per_rank, - args.hidden, args.intermediate_hidden, - args.num_experts, args.num_shared_experts, - args.num_topk, - args.config.block_m, args.config.block_n, args.config.block_k, - args.config.store_block_m, - args.config.num_ring_tokens, - args.config.num_stages, - args.config.num_bytes_per_pull, - 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"); - } - - 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.y, - args.cumulative_local_expert_recv_stats, - args.num_tokens, - args.sym_buffer_ptrs, - args.tensor_map_l1_acts, - args.tensor_map_l1_weights, - args.tensor_map_l1_output, - args.tensor_map_l2_acts, - args.tensor_map_l2_weights, - args.tensor_map_shared_l1_acts, - args.tensor_map_shared_l1_weights, - args.tensor_map_shared_l1_output, - args.tensor_map_shared_l2_acts, - args.tensor_map_shared_l2_weights - )); - } -}; - static void sm100_bf16_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_acts, const torch::Tensor& l2_acts, @@ -198,39 +103,67 @@ static void sm100_bf16_mega_moe( 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 SM100BF16MegaMoERuntime::Args args = { - .num_max_tokens_per_rank = num_max_tokens_per_rank, - .hidden = hidden, .intermediate_hidden = intermediate_hidden, - .num_experts = num_experts, .num_shared_experts = num_shared_experts, - .num_topk = num_topk, - .num_ranks = num_ranks, - .activation_clamp = activation_clamp, - .fast_math = fast_math, - .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_weights = tensor_map_l1_weights, - .tensor_map_l1_output = tensor_map_l1_output, - .tensor_map_l2_acts = tensor_map_l2_acts, - .tensor_map_l2_weights = tensor_map_l2_weights, - .tensor_map_shared_l1_acts = tensor_map_shared_l1_acts, - .tensor_map_shared_l1_weights = tensor_map_shared_l1_weights, - .tensor_map_shared_l1_output = tensor_map_shared_l1_output, - .tensor_map_shared_l2_acts = tensor_map_shared_l2_acts, - .tensor_map_shared_l2_weights = tensor_map_shared_l2_weights, - .launch_args = LaunchArgs(num_sms, - config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, - config.smem_size, 2) - }; + const auto num_sms = runtime->get_num_sms(); + + // Compile + const auto kernel = jit->compile("sm100_bf16_mega_moe", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_bf16_mega_moe_impl< + {}, + {}, {}, + {}, {}, + {}, {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {} + >); +}}; +)", num_max_tokens_per_rank, + hidden, intermediate_hidden, + num_experts, num_shared_experts, + num_topk, + config.block_m, config.block_n, config.block_k, + config.store_block_m, + config.num_ring_tokens, + config.num_stages, + config.num_bytes_per_pull, + config.num_dispatch_threads, config.num_non_epilogue_threads, config.num_epilogue_threads, + num_sms, num_ranks, + to_string(activation_clamp), + fast_math ? "true" : "false")); - const auto code = SM100BF16MegaMoERuntime::generate(args); - const auto runtime = compiler->build("sm100_bf16_mega_moe", code); - SM100BF16MegaMoERuntime::launch(runtime, args); + // Launch + jit->launch( + kernel, { + .num_smem_bytes = config.smem_size, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, 1, 1), + .cluster_dim = dim3(2, 1, 1), + }, + y.data_ptr(), + cumulative_local_expert_recv_stats_ptr, + num_tokens, + layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + tensor_map_l1_acts, + tensor_map_l1_weights, + tensor_map_l1_output, + tensor_map_l2_acts, + tensor_map_l2_weights, + tensor_map_shared_l1_acts, + tensor_map_shared_l1_weights, + tensor_map_shared_l1_output, + tensor_map_shared_l2_acts, + tensor_map_shared_l2_weights); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp index 65c9d501c2..e9a14c20f1 100644 --- a/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp @@ -1,65 +1,17 @@ #pragma once +#include +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../../utils/math.hpp" #include "../heuristics/sm100.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM100BmkBnkMnRuntime final: public LaunchRuntime { -public: - struct Args { - int s, m, n, k; - int block_m, block_n, block_k; - int split_factor; - int swizzle_ab_mode, swizzle_cd_mode; - int num_stages; - int num_threads; - - LaunchArgs launch_args; - - CUtensorMap tensor_map_a; - CUtensorMap tensor_map_b; - CUtensorMap tensor_map_d; - }; - - 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(&sm100_bmn_bnk_mn_gemm_impl< - {}, {}, {}, - {}, {}, {}, - {}, - {}, {}, - {}, {} - >); -}}; -)", - args.m, args.n, args.k, - args.block_m, args.block_n, args.block_k, - args.split_factor, - args.swizzle_ab_mode, args.swizzle_cd_mode, - args.num_stages, args.num_threads); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.s, args.tensor_map_a, args.tensor_map_b, args.tensor_map_d)); - } -}; - - static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a, const torch::Tensor &b, const torch::Tensor &d, @@ -76,7 +28,7 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a, const int swizzle_cd_mode = get_swizzle_mode(block_n, static_cast(d.element_size())); // Get best config - const int num_sms = device_runtime->get_num_sms(); + const int num_sms = runtime->get_num_sms(); const int num_mn_blocks = ceil_div(m, block_m) * ceil_div(n, block_n); const int num_sk_blocks = s * (k / block_k); const int split_factor = ceil_div(num_sk_blocks, std::max(num_sms / num_mn_blocks, 1)); @@ -104,7 +56,7 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a, DG_HOST_ASSERT(num_stages > 0); // Print configs - if (get_env("DG_JIT_DEBUG", 0)) { + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { printf("S: %d, M: %d, N: %d, K: %d -> " "block M: %d, block N: %d, block K: %d, split-K factor: %d" "stages: %d, shared memory: %d, swizzle AB: %d, swizzle CD: %d\n", @@ -116,22 +68,37 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a, const auto tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode); const auto tensor_map_d = make_tma_2d_desc(d, n, m, block_n, block_m, n, swizzle_cd_mode); - const SM100BmkBnkMnRuntime::Args args = { - .s = s, .m = m, .n = n, .k = k, - .block_m = block_m, .block_n = block_n, .block_k = block_k, - .split_factor = split_factor, - .swizzle_ab_mode = swizzle_ab_mode, - .swizzle_cd_mode = swizzle_cd_mode, - .num_stages = num_stages, - .num_threads = num_threads, - .launch_args = LaunchArgs(num_mn_blocks * ceil_div(num_sk_blocks, split_factor), num_threads, smem_size), - .tensor_map_a = tensor_map_a, - .tensor_map_b = tensor_map_b, - .tensor_map_d = tensor_map_d - }; - const auto code = SM100BmkBnkMnRuntime::generate(args); - const auto runtime = compiler->build("sm100_bmn_bnk_mn_gemm", code); - SM100BmkBnkMnRuntime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("sm100_bmn_bnk_mn_gemm", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_bmn_bnk_mn_gemm_impl< + {}, {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, {} + >); +}}; +)", + m, n, k, + block_m, block_n, block_k, + split_factor, + swizzle_ab_mode, swizzle_cd_mode, + num_stages, num_threads)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_mn_blocks * ceil_div(num_sk_blocks, split_factor), 1, 1), + .block_dim = dim3(num_threads, 1, 1), + }, + s, tensor_map_a, tensor_map_b, tensor_map_d + ); } } // namespace deep_gemm 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 9e4ba58b60..0e637da1ea 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp @@ -1,12 +1,10 @@ #pragma once +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../../utils/math.hpp" #include "../heuristics/sm100.hpp" @@ -15,14 +13,14 @@ namespace deep_gemm { -class SM100FP8FP4Gemm1D1DRuntime final: public LaunchRuntime { +class SM100FP8FP4Gemm1D1DRuntime final { public: struct Args { GemmDesc gemm_desc; GemmConfig gemm_config; - LaunchArgs launch_args; + deep_jit::cuda::LaunchOptions options; // TODO: move into descriptor - const std::optional epilogue_type; + EpilogueInput epilogue; // TODO: move into descriptor int gran_k_a, gran_k_b, k_alignment; @@ -35,28 +33,30 @@ class SM100FP8FP4Gemm1D1DRuntime final: public LaunchRuntimecompile(tag, std::format(R"( #include using namespace deep_gemm; static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm100_fp8_fp4_gemm_1d1d_impl< - {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, - {}, {}, {}, {}, {}, + {}, {}, {}, + {}, {}, {}, {}, {}, {}, - {}, - {}, {}, {}, {}, + {}, {}, {}, {}, {}, {} >); @@ -70,26 +70,35 @@ static void __instantiate_kernel() {{ args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k, args.gemm_desc.num_groups, args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode, - args.gemm_config.pipeline_config.num_stages, + args.gemm_config.pipeline_config.num_stages, args.gemm_config.pipeline_config.num_tma_store_stages, args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads, args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1, args.gemm_config.launch_config.num_sms, args.gemm_config.layout.swap_ab, args.gemm_desc.ensure_zero_padding, to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation, - to_string(args.gemm_desc.a_dtype), to_string(args.gemm_desc.b_dtype), to_string(args.gemm_desc.cd_dtype), - get_default_epilogue_type(args.epilogue_type)); - } + a_dtype, b_dtype, to_string(args.gemm_desc.cd_dtype), + args.epilogue.type)); - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - // TODO: optimize `args` copy - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + // Launch + jit->launch( + kernel, args.options, args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, + args.epilogue.args, args.tensor_map_a, args.tensor_map_b, args.tensor_map_sfa, args.tensor_map_sfb, - args.tensor_map_cd)); + args.tensor_map_cd + ); } }; +// Compute the SF block sizes for the SF TMA descriptors. +// Both MXF4 and MXF8F6F4 use UTCCP-aligned SF blocks, with a per-128 K sub-block count (`block_k / 128`). +static std::tuple get_sf_block_config(const GemmConfig& config, const GemmDesc& desc) { + const auto [sf_block_m, sf_block_n] = SM100ArchSpec::get_sf_uttcp_aligned_block_sizes( + config.layout.block_m, config.layout.block_n, desc.get_mma_kind()); + return {sf_block_m, sf_block_n, config.layout.block_k / 128}; +} + static void sm100_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, const torch::Tensor& b, const torch::Tensor& sfb, const std::optional& c, @@ -98,7 +107,8 @@ static void sm100_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& const int& gran_k_a, const int& gran_k_b, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const std::string& compiled_dims, - const std::optional& epilogue_type = std::nullopt) { + const std::optional& epilogue_type = std::nullopt, + const std::optional& alpha = std::nullopt) { const auto desc = GemmDesc { .gemm_type = GemmType::Normal, .kernel_type = KernelType::Kernel1D1D, @@ -107,41 +117,44 @@ static void sm100_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); - const auto cd = c.value_or(d); const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k, config.storage_config.load_block_m, config.layout.block_k, static_cast(a.stride(get_non_contiguous_dim(major_a))), 1, - config.storage_config.swizzle_a_mode); + config.storage_config.swizzle_a_mode, 0, false, not desc.is_mxf4_mma()); const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, config.storage_config.load_block_n, config.layout.block_k, static_cast(b.stride(get_non_contiguous_dim(major_b))), 1, - config.storage_config.swizzle_b_mode); + config.storage_config.swizzle_b_mode, 0, false, not desc.is_mxf4_mma()); const auto tensor_map_cd = make_tma_cd_desc(d, m, static_cast(d.size(-1)), config.storage_config.store_block_m, config.storage_config.store_block_n, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); + const auto [sf_block_mn_a, sf_block_mn_b, sf_block_k] = get_sf_block_config(config, desc); const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, - config.layout.block_m, gran_k_a, 1, 0); + sf_block_mn_a, gran_k_a, 1, 0, 0, false, sf_block_k); const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, - config.layout.block_n, gran_k_b, 1, 0); + sf_block_mn_b, gran_k_b, 1, 0, 0, false, sf_block_k); - // Launch - const SM100FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM100FP8FP4Gemm1D1DRuntime::compile_and_launch("sm100_fp8_fp4_gemm_1d1d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), - .epilogue_type = epilogue_type, + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, + .epilogue = make_epilogue_input(m, n, epilogue_type, alpha), .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, // NOTES: `k_alignment` is only used by k-grouped psum, dummy here @@ -152,10 +165,7 @@ static void sm100_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm100_fp8_fp4_gemm_1d1d", code); - SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); } static void sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -186,8 +196,8 @@ static void sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), .compiled_dims = compiled_dims, .ensure_zero_padding = ensure_zero_padding, .expected_m = expected_m_for_psum_layout.value_or(m), @@ -201,30 +211,33 @@ static void sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, config.storage_config.load_block_m, config.layout.block_k, static_cast(a.stride(get_non_contiguous_dim(major_a))), 1, - config.storage_config.swizzle_a_mode); + config.storage_config.swizzle_a_mode, 0, false, not desc.is_mxf4_mma()); const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, config.storage_config.load_block_n, config.layout.block_k, static_cast(b.stride(get_non_contiguous_dim(major_b))), num_groups, - config.storage_config.swizzle_b_mode); + config.storage_config.swizzle_b_mode, 0, false, not desc.is_mxf4_mma()); const auto tensor_map_cd = make_tma_cd_desc(d, m, n, config.storage_config.store_block_m, config.storage_config.store_block_n, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); + const auto [sf_block_mn_a, sf_block_mn_b, sf_block_k] = get_sf_block_config(config, desc); const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, - config.layout.block_m, gran_k_a, 1, 0); + sf_block_mn_a, gran_k_a, 1, 0, 0, false, sf_block_k); const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, - config.layout.block_n, gran_k_b, num_groups, 0); + sf_block_mn_b, gran_k_b, num_groups, 0, 0, false, sf_block_k); - // Launch kernel - const SM100FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM100FP8FP4Gemm1D1DRuntime::compile_and_launch("sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), - .epilogue_type = std::nullopt, + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, // NOTES: `k_alignment` is only used by k-grouped psum, dummy here @@ -235,10 +248,7 @@ static void sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d", code); - SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); } static void sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -258,8 +268,8 @@ static void sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), .compiled_dims = compiled_dims, .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups }; @@ -270,30 +280,33 @@ static void sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con config.storage_config.load_block_m, config.layout.block_k, static_cast(a.stride(get_non_contiguous_dim(major_a))), num_groups, - config.storage_config.swizzle_a_mode); + config.storage_config.swizzle_a_mode, 0, false, not desc.is_mxf4_mma()); const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k, config.storage_config.load_block_n, config.layout.block_k, static_cast(b.stride(get_non_contiguous_dim(major_b))), num_groups, - config.storage_config.swizzle_b_mode); + config.storage_config.swizzle_b_mode, 0, false, not desc.is_mxf4_mma()); const auto tensor_map_cd = make_tma_cd_desc(d, m, n, config.storage_config.store_block_m, config.storage_config.store_block_n, static_cast(d.stride(-2)), num_groups, config.storage_config.swizzle_cd_mode); + const auto [sf_block_mn_a, sf_block_mn_b, sf_block_k] = get_sf_block_config(config, desc); const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, - config.layout.block_m, gran_k_a, num_groups, 0); + sf_block_mn_a, gran_k_a, num_groups, 0, 0, false, sf_block_k); const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, - config.layout.block_n, gran_k_b, num_groups, 0); + sf_block_mn_b, gran_k_b, num_groups, 0, 0, false, sf_block_k); - // Launch kernel - const SM100FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM100FP8FP4Gemm1D1DRuntime::compile_and_launch("sm100_m_grouped_fp8_fp4_gemm_masked_1d1d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), - .epilogue_type = std::nullopt, + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, // NOTES: `k_alignment` is only used by k-grouped psum, dummy here @@ -304,10 +317,7 @@ static void sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm100_m_grouped_fp8_fp4_gemm_masked_1d1d", code); - SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); } static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -322,10 +332,6 @@ static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::T const std::string& compiled_dims, const bool& use_psum_layout) { DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN); - DG_HOST_ASSERT(gran_k == 32 or gran_k == 128); - DG_HOST_ASSERT(k_alignment % 32 == 0); - const int gran_k_a = gran_k; - const int gran_k_b = gran_k; const auto num_groups = static_cast(grouped_layout.numel()); const auto sum_k = static_cast(a.size(0)); const auto expected_k = ceil_div(sum_k, num_groups); @@ -338,13 +344,14 @@ static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::T .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), .compiled_dims = compiled_dims, // NOTES: expected_k is not used in SM100 get_best_config yet. .expected_m = m, .expected_n = n, .expected_k = expected_k, .expected_num_groups = num_groups }; const auto config = get_best_config(desc); + DG_HOST_ASSERT(k_alignment % config.layout.block_k == 0); // Create tensor descriptors const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k, @@ -357,26 +364,30 @@ static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::T config.layout.block_k, static_cast(b.stride(0)), 1, config.storage_config.swizzle_b_mode); - const auto tensor_map_cd = make_tma_cd_desc(d, m, n, - config.storage_config.store_block_m, + const auto tensor_map_cd = make_tma_3d_desc(d, n, m, num_groups, config.storage_config.store_block_n, - static_cast(d.stride(1)), num_groups, + config.storage_config.store_block_m, 1, + static_cast(d.stride(1)), + static_cast(d.stride(0)), config.storage_config.swizzle_cd_mode); - const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, static_cast(sfa.size(0)) * gran_k_a * 4, - config.layout.block_m, gran_k_a, 1, 0); - const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, static_cast(sfb.size(0)) * gran_k_b * 4, - config.layout.block_n, gran_k_b, 1, 0); - - // Launch kernel - const SM100FP8FP4Gemm1D1DRuntime::Args args = { + const auto [sf_block_mn_a, sf_block_mn_b, sf_block_k] = get_sf_block_config(config, desc); + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, static_cast(sfa.size(0)) * gran_k * 4, + sf_block_mn_a, gran_k, 1, 0, 0, false, sf_block_k); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, static_cast(sfb.size(0)) * gran_k * 4, + sf_block_mn_b, gran_k, 1, 0, 0, false, sf_block_k); + + // Compile and launch + SM100FP8FP4Gemm1D1DRuntime::compile_and_launch("sm100_k_grouped_fp8_gemm_1d1d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), - .epilogue_type = std::nullopt, - .gran_k_a = gran_k_a, - .gran_k_b = gran_k_b, + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, + .gran_k_a = gran_k, + .gran_k_b = gran_k, .k_alignment = k_alignment, .grouped_layout = grouped_layout.data_ptr(), .tensor_map_a = tensor_map_a, @@ -384,10 +395,7 @@ static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::T .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, .tensor_map_cd = tensor_map_cd - }; - const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm100_k_grouped_fp8_gemm_1d1d", code); - SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); } static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, @@ -397,7 +405,8 @@ static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, const int& batch_size, const int& m, const int& n, const int& k, const int& gran_k_a, const int& gran_k_b, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, - const std::string& compiled_dims) { + const std::string& compiled_dims, + const std::optional& sfd = std::nullopt) { const auto desc = GemmDesc { .gemm_type = GemmType::Batched, .kernel_type = KernelType::Kernel1D1D, @@ -406,10 +415,11 @@ static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), .compiled_dims = compiled_dims }; + DG_HOST_ASSERT(sfd.has_value() == (d.scalar_type() == torch::kFloat8_e4m3fn)); const auto config = get_best_config(desc); const int load_block_m = config.storage_config.load_block_m; @@ -437,19 +447,23 @@ static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, d.stride(1), d.stride(0), config.storage_config.swizzle_cd_mode); + const auto [sf_block_mn_a, sf_block_mn_b, sf_block_k] = get_sf_block_config(config, desc); const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, - config.layout.block_m, gran_k_a, batch_size, 0); + sf_block_mn_a, gran_k_a, batch_size, 0, 0, false, sf_block_k); const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k, - config.layout.block_n, gran_k_b, batch_size, 0); + sf_block_mn_b, gran_k_b, batch_size, 0, 0, false, sf_block_k); - // Launch - const SM100FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM100FP8FP4Gemm1D1DRuntime::compile_and_launch("sm100_fp8_gemm_1d1d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), - .epilogue_type = std::nullopt, + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, + .epilogue = make_epilogue_input(m, n, std::nullopt, std::nullopt, sfd), .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, // NOTES: `k_alignment` is only used by k-grouped psum, dummy here @@ -460,10 +474,85 @@ static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, .tensor_map_cd = tensor_map_cd + }); +} + +static void sm100_k_grouped_fp4_gemm_1d1d( + const torch::Tensor& a, const torch::Tensor& sfa, + const torch::Tensor& b, const torch::Tensor& sfb, + const std::optional& c, + const torch::Tensor& d, + const int& m, const int& n, const int& sum_k, + const torch::Tensor& grouped_layout, + const int& k_alignment, + const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, + const std::string& compiled_dims, + const bool& use_psum_layout +) { + constexpr int gran_k = 32; + DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); + const auto num_groups = static_cast(grouped_layout.numel()); + const auto expected_k = ceil_div(sum_k, num_groups); + + const auto desc = GemmDesc { + .gemm_type = use_psum_layout ? GemmType::KGroupedContiguousWithPsumLayout : GemmType::KGroupedContiguous, + .kernel_type = KernelType::Kernel1D1D, + .m = m, .n = n, .k = sum_k, .num_groups = num_groups, + .a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(), + .cd_dtype = d.scalar_type(), + .major_a = major_a, .major_b = major_b, + .with_accumulation = c.has_value(), + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, + .expected_m = m, .expected_n = n, .expected_k = expected_k, .expected_num_groups = num_groups }; - const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm100_fp8_gemm_1d1d", code); - SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args); + const auto config = get_best_config(desc); + DG_HOST_ASSERT(k_alignment % config.layout.block_k == 0); + + // Create tensor descriptors + const auto tensor_map_a = make_tma_a_desc(major_a, a, m, sum_k, + config.storage_config.load_block_m, + config.layout.block_k, + static_cast(a.stride(0)), 1, + config.storage_config.swizzle_a_mode, 0, false, false); + const auto tensor_map_b = make_tma_b_desc(major_b, b, n, sum_k, + config.storage_config.load_block_n, + config.layout.block_k, + static_cast(b.stride(0)), 1, + config.storage_config.swizzle_b_mode, 0, false, false); + const auto tensor_map_cd = make_tma_3d_desc(d, n, m, num_groups, + config.storage_config.store_block_n, + config.storage_config.store_block_m, 1, + static_cast(d.stride(1)), + static_cast(d.stride(0)), + config.storage_config.swizzle_cd_mode); + const auto [sf_block_mn_a, sf_block_mn_b, sf_block_k] = get_sf_block_config(config, desc); + const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, static_cast(sfa.size(0)) * gran_k * 4, + sf_block_mn_a, gran_k, 1, 0, 0, false, sf_block_k); + const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, static_cast(sfb.size(0)) * gran_k * 4, + sf_block_mn_b, gran_k, 1, 0, 0, false, sf_block_k); + + // Compile and launch + SM100FP8FP4Gemm1D1DRuntime::compile_and_launch("sm100_k_grouped_fp4_gemm_1d1d", { + .gemm_desc = desc, + .gemm_config = config, + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, + .gran_k_a = gran_k, + .gran_k_b = gran_k, + .k_alignment = k_alignment, + .grouped_layout = grouped_layout.data_ptr(), + .tensor_map_a = tensor_map_a, + .tensor_map_b = tensor_map_b, + .tensor_map_sfa = tensor_map_sfa, + .tensor_map_sfb = tensor_map_sfb, + .tensor_map_cd = tensor_map_cd + }); } } // namespace deep_gemm 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 9c4be4b08d..562f3042e9 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -1,133 +1,18 @@ #pragma once +#include #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 "../../runtime/runtime.hpp" +#include "../../utils/exception.hpp" #include "../heuristics/mega_moe.hpp" +#include "runtime_utils.hpp" namespace deep_gemm { -class SM100FP8FP4MegaMoERuntime final : public LaunchRuntime { -public: - struct Args { - // Templated arguments - int num_max_tokens_per_rank; - int hidden, intermediate_hidden; - int num_experts, num_shared_experts, num_topk; - int num_ranks; - float activation_clamp; - bool fast_math; - MegaMoEConfig config; - - // Runtime arguments - void* y; - int* cumulative_local_expert_recv_stats; - int num_tokens; - layout::SymBuffer<> sym_buffer_ptrs; - - // Tensormap - CUtensorMap tensor_map_l1_acts; - CUtensorMap tensor_map_l1_acts_sf; - CUtensorMap tensor_map_l1_weights; - CUtensorMap tensor_map_l1_weights_sf; - CUtensorMap tensor_map_l1_output; - CUtensorMap tensor_map_l2_acts; - CUtensorMap tensor_map_l2_acts_sf; - CUtensorMap tensor_map_l2_weights; - CUtensorMap tensor_map_l2_weights_sf; - CUtensorMap tensor_map_shared_l1_acts; - CUtensorMap tensor_map_shared_l1_acts_sf; - CUtensorMap tensor_map_shared_l1_weights; - CUtensorMap tensor_map_shared_l1_weights_sf; - CUtensorMap tensor_map_shared_l1_output; - CUtensorMap tensor_map_shared_l2_acts; - CUtensorMap tensor_map_shared_l2_acts_sf; - CUtensorMap tensor_map_shared_l2_weights; - CUtensorMap tensor_map_shared_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(&sm100_fp8_fp4_mega_moe_impl< - {}, - {}, {}, - {}, {}, - {}, {}, {}, - {}, - {}, {}, - {}, - {}, - {}, - {}, - {}, - {}, {}, {}, - {}, {}, - {}, - {} - >); -}}; -)", args.num_max_tokens_per_rank, - args.hidden, args.intermediate_hidden, - args.num_experts, args.num_shared_experts, - args.num_topk, - args.config.block_m, args.config.block_n, args.config.block_k, - args.config.store_block_m, - args.config.sf_block_m, args.config.sf_block_n, - args.config.num_ring_tokens, - args.config.num_sf_ring_tokens, - args.config.num_stages, - args.config.num_bytes_per_pull, - 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"); - } - - 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.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.tensor_map_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.tensor_map_l2_weights_sf, - args.tensor_map_shared_l1_acts, - args.tensor_map_shared_l1_acts_sf, - args.tensor_map_shared_l1_weights, - args.tensor_map_shared_l1_weights_sf, - args.tensor_map_shared_l1_output, - args.tensor_map_shared_l2_acts, - args.tensor_map_shared_l2_acts_sf, - args.tensor_map_shared_l2_weights, - args.tensor_map_shared_l2_weights_sf - )); - } -}; - static void sm100_fp8_fp4_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf, @@ -183,7 +68,7 @@ static void sm100_fp8_fp4_mega_moe( intermediate_hidden * 2, hidden, config.block_n, kGranK, num_experts_per_rank, 0, 0, false, - sf_smem_outer_dim); + sf_smem_outer_dim); // 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). @@ -273,47 +158,81 @@ static void sm100_fp8_fp4_mega_moe( 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 SM100FP8FP4MegaMoERuntime::Args args = { - .num_max_tokens_per_rank = num_max_tokens_per_rank, - .hidden = hidden, .intermediate_hidden = intermediate_hidden, - .num_experts = num_experts, .num_shared_experts = num_shared_experts, - .num_topk = num_topk, - .num_ranks = num_ranks, - .activation_clamp = activation_clamp, - .fast_math = fast_math, - .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, - .tensor_map_l1_weights_sf = tensor_map_l1_weights_sf, - .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, - .tensor_map_l2_weights_sf = tensor_map_l2_weights_sf, - .tensor_map_shared_l1_acts = tensor_map_shared_l1_acts, - .tensor_map_shared_l1_acts_sf = tensor_map_shared_l1_acts_sf, - .tensor_map_shared_l1_weights = tensor_map_shared_l1_weights, - .tensor_map_shared_l1_weights_sf = tensor_map_shared_l1_weights_sf, - .tensor_map_shared_l1_output = tensor_map_shared_l1_output, - .tensor_map_shared_l2_acts = tensor_map_shared_l2_acts, - .tensor_map_shared_l2_acts_sf = tensor_map_shared_l2_acts_sf, - .tensor_map_shared_l2_weights = tensor_map_shared_l2_weights, - .tensor_map_shared_l2_weights_sf = tensor_map_shared_l2_weights_sf, - .launch_args = LaunchArgs(num_sms, - config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, - config.smem_size, 2) - }; + const auto num_sms = runtime->get_num_sms(); + + // Compile + const auto kernel = jit->compile("sm100_fp8_fp4_mega_moe", std::format(R"( +#include + +using namespace deep_gemm; - const auto code = SM100FP8FP4MegaMoERuntime::generate(args); - const auto runtime = compiler->build("sm100_fp8_fp4_mega_moe", code); - SM100FP8FP4MegaMoERuntime::launch(runtime, args); +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_fp8_fp4_mega_moe_impl< + {}, + {}, {}, + {}, {}, + {}, {}, {}, + {}, + {}, {}, + {}, + {}, + {}, + {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {}, + {} + >); +}}; +)", num_max_tokens_per_rank, + hidden, intermediate_hidden, + num_experts, num_shared_experts, + num_topk, + config.block_m, config.block_n, config.block_k, + config.store_block_m, + config.sf_block_m, config.sf_block_n, + config.num_ring_tokens, + config.num_sf_ring_tokens, + config.num_stages, + config.num_bytes_per_pull, + config.num_dispatch_threads, config.num_non_epilogue_threads, config.num_epilogue_threads, + num_sms, num_ranks, + to_string(activation_clamp), + fast_math ? "true" : "false", + to_string(l1_weights.scalar_type()))); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = config.smem_size, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads, 1, 1), + .cluster_dim = dim3(2, 1, 1), + }, + y.data_ptr(), + cumulative_local_expert_recv_stats_ptr, + num_tokens, + layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), + tensor_map_l1_acts, + tensor_map_l1_acts_sf, + tensor_map_l1_weights, + tensor_map_l1_weights_sf, + tensor_map_l1_output, + tensor_map_l2_acts, + tensor_map_l2_acts_sf, + tensor_map_l2_weights, + tensor_map_l2_weights_sf, + tensor_map_shared_l1_acts, + tensor_map_shared_l1_acts_sf, + tensor_map_shared_l1_weights, + tensor_map_shared_l1_weights_sf, + tensor_map_shared_l1_output, + tensor_map_shared_l2_acts, + tensor_map_shared_l2_acts_sf, + tensor_map_shared_l2_weights, + tensor_map_shared_l2_weights_sf); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_mega_mhc.hpp b/csrc/jit_kernels/impls/sm100_mega_mhc.hpp new file mode 100644 index 0000000000..3be09f3fa6 --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_mega_mhc.hpp @@ -0,0 +1,149 @@ +#pragma once + +#include +#include + +#include + +#include "../../runtime/jit.hpp" +#include "../heuristics/sm100.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +namespace mhc_layout = layout::mega_mhc; + +static void sm100_mega_mhc( + const torch::Tensor& x, const torch::Tensor& residual, + const torch::Tensor& post_mix, const torch::Tensor& comb_res_mix, + const std::optional& shifted_prev_mix, const torch::Tensor& fn, + const torch::Tensor& mix_scales, const torch::Tensor& mix_bases, + const float& hc_norm_eps, const float& hc_pre_eps, const float& hc_post_scale, + const float& sinkhorn_eps, const int& num_sinkhorn_iters, + const torch::Tensor& rmsnorm_weight, const float& rmsnorm_eps, + const torch::Tensor& new_residual, + const std::optional& new_prev_mix, const torch::Tensor& new_post_mix, + const torch::Tensor& new_comb_res_mix, const torch::Tensor& y_bf16, + const bool& store_bf16, + const torch::Tensor& scratch, + const torch::Tensor& split_barriers, + const std::optional& y_fp8, + const std::optional& y_primary_sf, + const std::optional& y_shared_sf, + const int& shared_sf_block_m, const float& rmsnorm_scale, + const int& num_tokens, const int& hidden, const int& num_splits, const int& num_sms) { + const bool is_shifted = shifted_prev_mix.has_value(); + + const auto make_hidden_tma_desc = [&](const torch::Tensor& tensor) { + return make_tma_2d_desc(tensor, hidden, num_tokens, mhc_layout::BLOCK_K, mhc_layout::BLOCK_M, + static_cast(tensor.stride(0)), mhc_layout::kSwizzleMode); + }; + const auto make_residual_tma_desc = [&](const torch::Tensor& tensor) { + return make_tma_3d_desc( + tensor, hidden, num_tokens, mhc_layout::kNumRoutes, + mhc_layout::BLOCK_K, mhc_layout::BLOCK_M, mhc_layout::kNumRoutes, + static_cast(tensor.stride(0)), static_cast(tensor.stride(1)), mhc_layout::kSwizzleMode); + }; + const auto make_coeff_tma_desc = [&](const torch::Tensor& tensor, const int channels, const int swizzle_mode) { + return make_tma_2d_desc( + tensor, channels, num_tokens, channels, mhc_layout::BLOCK_M, + static_cast(tensor.stride(0)), swizzle_mode); + }; + + const auto smem_size = static_cast(sizeof(mhc_layout::SharedStorage)); + DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity); + + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { + printf("Mega %s mHC: T=%d, H=%d, splits=%d, IO/Fn/A stages=%u/%u/%u, " + "threads=%u, shared memory=%d\n", + is_shifted ? "shifted" : "normal", num_tokens, hidden, num_splits, + mhc_layout::kNumIOStages, mhc_layout::kNumFnStages, mhc_layout::kNumAStages, + mhc_layout::kNumThreads, smem_size); + } + + // Tensor maps use fastest-moving dimensions first. Residual and Fn expose + // the HC route as a third dimension without materializing transposed views. + const auto tensor_map_residual = make_residual_tma_desc(residual); + const auto tensor_map_x = make_hidden_tma_desc(x); + const auto tensor_map_fn = make_tma_3d_desc( + fn, + hidden, mhc_layout::kNumHCOutputs, mhc_layout::kNumRoutes, + mhc_layout::BLOCK_K, mhc_layout::kNumHCOutputs, mhc_layout::kNumRoutes, + static_cast(fn.stride(0)), hidden, + mhc_layout::kSwizzleMode, 0, true); + const auto tensor_map_post_mix = make_coeff_tma_desc(post_mix, mhc_layout::kNumRoutes, 0); + const auto tensor_map_comb_res_mix = make_coeff_tma_desc( + comb_res_mix, mhc_layout::kNumRoutes * mhc_layout::kNumRoutes, mhc_layout::kSwizzleMode / 2); + const auto tensor_map_shifted_prev_mix = make_coeff_tma_desc( + is_shifted ? shifted_prev_mix.value() : post_mix, mhc_layout::kNumRoutes, 0); + const auto tensor_map_new_residual = make_residual_tma_desc(new_residual); + const auto tensor_map_y_bf16 = make_hidden_tma_desc(is_shifted ? y_bf16 : x); + + const mhc_layout::MixArgs mix_args = { + .scales = mix_scales.data_ptr(), + .bases = mix_bases.data_ptr(), + .new_prev_mix = is_shifted ? new_prev_mix->data_ptr() : nullptr, + .new_post_mix = new_post_mix.data_ptr(), + .new_comb_res_mix = new_comb_res_mix.data_ptr(), + .hc_norm_eps = hc_norm_eps, + .hc_pre_eps = hc_pre_eps, + .hc_post_scale = hc_post_scale, + .sinkhorn_eps = sinkhorn_eps, + .num_sinkhorn_iters = static_cast(num_sinkhorn_iters), + }; + const mhc_layout::NormArgs norm_args = { + .num_tokens = static_cast(num_tokens), + .weight = reinterpret_cast(rmsnorm_weight.data_ptr()), + .new_residual = reinterpret_cast(new_residual.data_ptr()), + .eps = rmsnorm_eps, + .scale = rmsnorm_scale, + .y_bf16 = reinterpret_cast(y_bf16.data_ptr()), + .y_fp8 = y_fp8.has_value() ? reinterpret_cast<__nv_fp8_e4m3*>(y_fp8->data_ptr()) : nullptr, + .y_primary_sf = y_primary_sf.has_value() ? reinterpret_cast(y_primary_sf->data_ptr()) : nullptr, + .y_primary_sf_stride_token = y_primary_sf.has_value() ? y_primary_sf->stride(0) : 0, + .y_primary_sf_stride_word = y_primary_sf.has_value() ? y_primary_sf->stride(1) : 0, + .y_shared_sf = y_shared_sf.has_value() ? reinterpret_cast(y_shared_sf->data_ptr()) : nullptr, + .y_shared_sf_stride_word = y_shared_sf.has_value() ? y_shared_sf->stride(1) : 0, + }; + + // Compile + const auto kernel = jit->compile("sm100_mega_mhc", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_mega_mhc_impl< + {}, {}, {}, + {}, {}, {}, {} + >); +}}; +)", + hidden, num_splits, num_sms, + mix_args.new_prev_mix != nullptr, + store_bf16, + norm_args.y_fp8 != nullptr, + shared_sf_block_m)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(mhc_layout::kNumThreads, 1, 1), + }, + tensor_map_residual, + tensor_map_x, + tensor_map_fn, + tensor_map_post_mix, + tensor_map_comb_res_mix, + tensor_map_shifted_prev_mix, + tensor_map_new_residual, + tensor_map_y_bf16, + mix_args, + norm_args, + scratch.data_ptr(), + reinterpret_cast(split_barriers.data_ptr())); +} + +} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_mqa_logits.hpp b/csrc/jit_kernels/impls/sm100_mqa_logits.hpp index c594d2bd6e..e9b68a2e53 100644 --- a/csrc/jit_kernels/impls/sm100_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/sm100_mqa_logits.hpp @@ -1,57 +1,16 @@ #pragma once -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include + +#include + +#include "../../runtime/runtime.hpp" #include "../heuristics/sm100.hpp" #include "runtime_utils.hpp" -#include namespace deep_gemm { // SM100 paged metadata emits per-SM starts as (q_token_idx, kv_split_idx) -class SM100PagedMQALogitsMetadataRuntime final: public LaunchRuntime { -public: - struct Args { - int next_n; - bool is_context_lens_2d; - bool is_varlen; - int split_kv; - int num_sms; - - int num_requests; - int num_q_tokens_total; - int* context_lens; - int* indices; - int* schedule_meta; - - LaunchArgs launch_args; - }; - - static std::string generate_impl(const Args& args) { - // Metadata weights by q-token count, so BLOCK_Q is a placeholder here - return fmt::format(R"( -#include - -using namespace deep_gemm; - -static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sched::sm100_paged_mqa_logits_metadata< - {}, {}, {}, 1, {}, {} - >); -}}; -)", args.next_n, args.is_context_lens_2d ? "true" : "false", args.is_varlen ? "true" : "false", - args.split_kv, args.num_sms); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.num_requests, args.num_q_tokens_total, - args.context_lens, args.indices, args.schedule_meta - )); - } -}; - static void sm100_paged_mqa_logits_metadata(const torch::Tensor& context_lens, const torch::Tensor& schedule_meta, const int& num_requests, const int& num_q_tokens_total, @@ -59,152 +18,133 @@ static void sm100_paged_mqa_logits_metadata(const torch::Tensor& context_lens, const bool& is_context_lens_2d, const bool& is_varlen, const int* indices_ptr) { constexpr int split_kv = 256; - const int num_threads = 256; - // smem: prefix_work[num_requests] + request_q_token_start[num_requests] - const int smem_size = 2 * num_requests * static_cast(sizeof(int)); + constexpr int num_threads = 1024; + // Request starts (varlen only), work prefix, warp sums, and block total + const int smem_size = ((is_varlen ? 2 : 1) * num_requests + num_threads / 32 + 1) * static_cast(sizeof(int)); DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity); - const SM100PagedMQALogitsMetadataRuntime::Args args = { - .next_n = next_n, - .is_context_lens_2d = is_context_lens_2d, - .is_varlen = is_varlen, - .split_kv = split_kv, - .num_sms = num_sms, - .num_requests = num_requests, - .num_q_tokens_total = num_q_tokens_total, - .context_lens = context_lens.data_ptr(), - .indices = const_cast(indices_ptr), - .schedule_meta = schedule_meta.data_ptr(), - .launch_args = LaunchArgs(1, num_threads, smem_size) - }; - const auto code = SM100PagedMQALogitsMetadataRuntime::generate(args); - const auto runtime = compiler->build("sm100_paged_mqa_logits_metadata", code); - SM100PagedMQALogitsMetadataRuntime::launch(runtime, args); -} + // Compile + const auto kernel = jit->compile("sm100_paged_mqa_logits_metadata", std::format(R"( +#include -// Sizes `MQALogitsSharedStorage` for a runtime (is_mx_sf, qk_dtype, num_heads, head_dim) -// BLOCK_Q = 128 / num_heads keeps UMMA_N = 128 for all supported head counts -static int get_mqa_logits_smem_size(const int& num_heads, const int& head_dim, - const bool& is_mx_sf, const at::ScalarType& qk_dtype) { - const auto get_smem_size = [&](auto is_mx_sf_c, auto h, auto d) { - constexpr bool kIsMXSF = decltype(is_mx_sf_c)::value; - constexpr int H = decltype(h)::value, D = decltype(d)::value; - constexpr int kNumKVStages = cute::is_same_v ? 10 : 5; - return static_cast(sizeof(layout::MQALogitsSharedStorage)); - }; - const auto dispatch_heads = [&](auto&& fn) { - switch (num_heads) { - case 8: return fn(cute::C<8>{}); - case 16: return fn(cute::C<16>{}); - case 32: return fn(cute::C<32>{}); - case 64: return fn(cute::C<64>{}); - default: DG_HOST_UNREACHABLE("Unsupported num_heads for MQA logits"); - } - }; - const auto dispatch_head_dim = [&](auto&& fn) { - switch (head_dim) { - case 32: return fn(cute::C<32>{}); - case 64: return fn(cute::C<64>{}); - case 128: return fn(cute::C<128>{}); - default: DG_HOST_UNREACHABLE("Unsupported head_dim for MQA logits"); - } - }; - const auto compute_smem_size = [&](auto is_mx_sf_c) { - return dispatch_heads([&](auto h) { - return dispatch_head_dim([&](auto d) { - switch (qk_dtype) { - case torch::kFloat8_e4m3fn: return get_smem_size.template operator()(is_mx_sf_c, h, d); - case kPackedFP4: return get_smem_size.template operator()(is_mx_sf_c, h, d); - default: DG_HOST_UNREACHABLE("Unsupported MQA logits QK dtype"); - } - }); - }); - }; +using namespace deep_gemm; - const int smem_size = is_mx_sf ? compute_smem_size(cute::C{}) : compute_smem_size(cute::C{}); - DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity); - return smem_size; +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sched::sm100_paged_mqa_logits_metadata< + {}, {}, {}, {}, {} + >); +}}; +)", next_n, is_context_lens_2d ? "true" : "false", is_varlen ? "true" : "false", + split_kv, num_sms)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(1, 1, 1), + .block_dim = dim3(num_threads, 1, 1), + }, + num_requests, num_q_tokens_total, + context_lens.data_ptr(), const_cast(indices_ptr), schedule_meta.data_ptr() + ); } -// Unified contiguous-KV runtime for FP8 / MXFP4 / MXFP8; FP8 reuses the unused `sf_q` descriptor slot -class SM100MQALogitsRuntime final: public LaunchRuntime { -public: - struct Args { - int num_q_tokens; - int num_kv_tokens; - int stride_logits; - int num_heads, head_dim; - bool is_mx_sf, is_compressed_logits; - - int num_q_stages; - int num_kv_stages; - int block_q; - int split_kv; - - int* cu_seq_len_k_start; - int* cu_seq_len_k_end; - void* logits; - - CUtensorMap tensor_map_q; - CUtensorMap tensor_map_sf_q; - CUtensorMap tensor_map_kv; - CUtensorMap tensor_map_sf_kv; - CUtensorMap tensor_map_weights; - at::ScalarType qk_dtype; - at::ScalarType logits_dtype; - at::ScalarType weights_dtype; - - int num_specialized_threads; - int num_math_threads; - - LaunchArgs launch_args; - }; +// Costs in units of one K split of main-loop work +static constexpr int kMQALogitsSplitCost = 1; +static constexpr int kMQALogitsSegmentCost = 1; +// Unroll depth of the builder's boundary search +static constexpr int kMQALogitsMaxNumQBlocksLog2 = 15; + +static int get_mqa_logits_metadata_num_words(const int& num_q_tokens, const int& block_q, const int& num_sms) { + DG_HOST_ASSERT(num_q_tokens > 0 and block_q > 0 and num_sms > 0); + const int num_q_blocks = math::ceil_div(num_q_tokens, block_q); + DG_HOST_ASSERT(num_q_blocks <= (1 << kMQALogitsMaxNumQBlocksLog2)); + return (3 * num_sms + 1) / 2 * 2 + 2 * num_q_blocks; +} - static std::string generate_impl(const Args& args) { - DG_HOST_ASSERT(128 % args.num_heads == 0); +static void sm100_mqa_logits_metadata(const torch::Tensor& cu_seq_len_k_start, + const torch::Tensor& cu_seq_len_k_end, + const torch::Tensor& schedule_meta, + const int& num_q_tokens, + const int& num_kv_tokens, + const int& block_q, + const int& split_kv, + const int& num_sms) { + DG_STATIC_ASSERT(kMQALogitsSegmentCost > 0 and kMQALogitsSplitCost > 0, "Invalid schedule cost constants"); + const int num_q_blocks = math::ceil_div(num_q_tokens, block_q); + const int64_t total_work_bound = static_cast(num_q_blocks) * math::ceil_div(num_kv_tokens, split_kv); + const int64_t total_cost_bound = kMQALogitsSplitCost * total_work_bound + kMQALogitsSegmentCost * num_q_blocks; + DG_HOST_ASSERT(total_work_bound <= std::numeric_limits::max()); + DG_HOST_ASSERT(total_cost_bound <= std::numeric_limits::max()); + constexpr int num_threads = 256; + // The builder keeps the per-block work and cost prefixes in shared memory + const int smem_size = 2 * num_q_blocks * static_cast(sizeof(uint32_t)); + DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity); - return fmt::format(R"( -#include + // Compile + const auto kernel = jit->compile("sm100_mqa_logits_metadata", std::format(R"( +#include using namespace deep_gemm; static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sm100_mqa_logits< - {}, {}, - {}, - {}, {}, - {}, {}, - {}, - {}, {}, - {}, {}, - {}, {} + auto ptr = reinterpret_cast(&sched::sm100_mqa_logits_metadata< + {}, {}, {}, {}, {}, {} >); }}; -)", args.num_heads, args.head_dim, - args.is_mx_sf ? "true" : "false", - args.is_compressed_logits, - args.block_q, args.split_kv, - args.num_q_stages, args.num_kv_stages, - args.launch_args.grid_dim.first, - args.num_specialized_threads, args.num_math_threads, - args.qk_dtype == kPackedFP4 ? "cutlass::float_e2m1_t" : "cutlass::float_e4m3_t", - to_string(args.logits_dtype), - args.weights_dtype == torch::kBFloat16 ? "__nv_bfloat16" : "float"); - } +)", block_q, split_kv, num_sms, + kMQALogitsSegmentCost, kMQALogitsSplitCost, kMQALogitsMaxNumQBlocksLog2)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(1, 1, 1), + .block_dim = dim3(num_threads, 1, 1), + }, + num_q_tokens, num_kv_tokens, + cu_seq_len_k_start.data_ptr(), cu_seq_len_k_end.data_ptr(), + schedule_meta.data_ptr() + ); +} - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.num_q_tokens, args.num_kv_tokens, - args.stride_logits, - args.cu_seq_len_k_start, args.cu_seq_len_k_end, - args.logits, - args.tensor_map_q, args.tensor_map_sf_q, - args.tensor_map_kv, args.tensor_map_sf_kv, - args.tensor_map_weights - )); - } -}; +static int get_mqa_logits_smem_size(const int& num_heads, const int& head_dim, + const bool& is_mx_sf, const at::ScalarType& qk_dtype, + const at::ScalarType& weights_dtype, + const int& block_q, const int& split_kv, + const int& num_q_stages, const int& num_kv_stages, + const int& num_tmem_stages) { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + constexpr uint32_t kNumUTCCPAlignedElems = 128; + constexpr uint32_t kTmaAlignment = 128; + const bool is_fp4 = qk_dtype == kPackedFP4; + const uint32_t num_qk_bytes_per_token = is_fp4 ? head_dim / 2 : head_dim; + const uint32_t swizzle_alignment = 8 * num_qk_bytes_per_token; + const uint32_t block_qh = block_q * num_heads; + const uint32_t umma_n = math::align(block_qh, 8u); + const uint32_t num_sf_q = is_mx_sf ? math::align(umma_n, kNumUTCCPAlignedElems) : 1; + const uint32_t num_sf_kv = is_mx_sf ? math::align(static_cast(split_kv), kNumUTCCPAlignedElems) : split_kv; + const uint32_t num_reduce_bytes = weights_dtype == torch::kBFloat16 ? sizeof(nv_bfloat16) : sizeof(float); + + uint32_t num_smem_bytes = 0; + const auto add_region = [&](const uint32_t& num_region_bytes, const uint32_t& alignment) { + num_smem_bytes = math::align(num_smem_bytes, alignment) + num_region_bytes; + }; + add_region(num_q_stages * umma_n * num_qk_bytes_per_token, swizzle_alignment); + add_region(num_kv_stages * split_kv * num_qk_bytes_per_token, swizzle_alignment); + add_region(num_q_stages * num_sf_q * sizeof(uint32_t), kTmaAlignment); + add_region(num_kv_stages * num_sf_kv * sizeof(uint32_t), kTmaAlignment); + const uint32_t num_weight_bytes_per_row = math::align(num_heads * num_reduce_bytes, 16u); + const uint32_t num_weight_bytes_per_stage = math::align(block_q * num_weight_bytes_per_row, kTmaAlignment); + add_region(num_q_stages * num_weight_bytes_per_stage, kTmaAlignment); + const uint32_t num_barriers = 2 * (num_q_stages + num_kv_stages + num_tmem_stages); + add_region(num_barriers * sizeof(Barrier), alignof(Barrier)); + add_region(sizeof(uint32_t), alignof(uint32_t)); + return static_cast(math::align(num_smem_bytes, swizzle_alignment)); +} +// Unified contiguous-KV runtime for FP8 / MXFP4 / MXFP8; FP8 reuses the unused `sf_q` descriptor slot static void sm100_mqa_logits(const torch::Tensor& q, const std::optional& sf_q, const torch::Tensor& kv, const torch::Tensor& sf_kv, const torch::Tensor& weights, @@ -216,8 +156,10 @@ static void sm100_mqa_logits(const torch::Tensor& q, const std::optional& schedule_meta) { const bool is_fp4 = qk_dtype == kPackedFP4; constexpr int num_specialized_threads = 128; @@ -227,7 +169,20 @@ static void sm100_mqa_logits(const torch::Tensor& q, const std::optional 0); - + const int num_sms = runtime->get_num_sms(); + DG_HOST_ASSERT(not (clean_logits and is_compressed_logits)); + + if (schedule_meta.has_value()) { + const auto& workspace = schedule_meta.value(); + DG_HOST_ASSERT(workspace.is_cuda()); + DG_HOST_ASSERT(workspace.device() == q.device()); + DG_HOST_ASSERT(workspace.scalar_type() == torch::kInt32); + DG_HOST_ASSERT(workspace.is_contiguous()); + // Metadata is accessed as 8-byte words + DG_HOST_ASSERT(reinterpret_cast(workspace.data_ptr()) % 8 == 0); + const int required_words = get_mqa_logits_metadata_num_words(num_q_tokens, block_q, num_sms); + DG_HOST_ASSERT(workspace.numel() >= required_words); + } // MX SF formats consume `sf_q`; FP8 fills that descriptor slot with KV scales CUtensorMap tensor_map_q, tensor_map_sf_q, tensor_map_kv, tensor_map_sf_kv; if (is_fp4) @@ -254,130 +209,81 @@ static void sm100_mqa_logits(const torch::Tensor& q, const std::optional(weights.stride(0)), 0); - - const int smem_size = get_mqa_logits_smem_size(num_heads, head_dim, is_mx_sf, qk_dtype); - - const SM100MQALogitsRuntime::Args args = { - .num_q_tokens = num_q_tokens, - .num_kv_tokens = num_kv_tokens, - .stride_logits = stride_logits, - .num_heads = num_heads, .head_dim = head_dim, - .is_mx_sf = is_mx_sf, - .is_compressed_logits = is_compressed_logits, - .num_q_stages = num_q_stages, - .num_kv_stages = num_kv_stages, - .block_q = block_q, - .split_kv = split_kv, - .cu_seq_len_k_start = cu_seq_len_k_start.data_ptr(), - .cu_seq_len_k_end = cu_seq_len_k_end.data_ptr(), - .logits = logits.data_ptr(), - .tensor_map_q = tensor_map_q, - .tensor_map_sf_q = tensor_map_sf_q, - .tensor_map_kv = tensor_map_kv, - .tensor_map_sf_kv = tensor_map_sf_kv, - .tensor_map_weights = tensor_map_weights, - .qk_dtype = qk_dtype, - .logits_dtype = logits_dtype, - .weights_dtype = weights.scalar_type(), - .num_specialized_threads = num_specialized_threads, - .num_math_threads = num_math_threads, - .launch_args = LaunchArgs(device_runtime->get_num_sms(), - num_specialized_threads + num_math_threads, - smem_size) - }; - const auto code = SM100MQALogitsRuntime::generate(args); - const auto runtime = compiler->build("sm100_mqa_logits", code); - SM100MQALogitsRuntime::launch(runtime, args); -} - -// Paged variant: separate host/runtime path, shared device core - -// Unified paged runtime for FP8 / MXFP4 / MXFP8; FP8 reuses the unused `sf_q` descriptor slot -class SM100PagedMQALogitsRuntime final: public LaunchRuntime { -public: - struct Args { - int num_q_tokens_total; - int tokens_per_request; - int num_heads; - int head_dim; - int page_kv; - bool is_mx_sf; - bool is_context_lens_2d; - bool is_varlen; - int block_table_stride; - int logits_stride; - - int num_q_stages; - int num_kv_stages; - int split_kv; - int splits_per_chunk; - - int* context_lens; - void* logits; - int* block_table; - int* indices; - int* schedule_meta; - - CUtensorMap tensor_map_q; - CUtensorMap tensor_map_sf_q; - CUtensorMap tensor_map_kv; - CUtensorMap tensor_map_sf_kv; - CUtensorMap tensor_map_weights; - at::ScalarType qk_dtype; - at::ScalarType logits_dtype; - at::ScalarType weights_dtype; - - int num_specialized_threads; - int num_math_threads; - - LaunchArgs launch_args; - }; + const uint32_t weights_stride = static_cast(weights.stride(0)); + const uint32_t num_weight_row_bytes = num_heads * static_cast(weights.element_size()); + const uint32_t num_weight_elements_per_row = math::align(num_weight_row_bytes, 16u) / static_cast(weights.element_size()); + const auto tensor_map_weights = make_tma_2d_desc( + weights, num_heads, num_q_tokens, + static_cast(num_weight_elements_per_row), block_q, + static_cast(weights_stride), 0); + + const int smem_size = get_mqa_logits_smem_size( + num_heads, head_dim, is_mx_sf, qk_dtype, weights.scalar_type(), + block_q, split_kv, num_q_stages, num_kv_stages, 3); + DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity); - static std::string generate_impl(const Args& args) { - return fmt::format(R"( + // Compile + const auto kernel = jit->compile("sm100_mqa_logits", std::format(R"( #include using namespace deep_gemm; +static_assert(sizeof(layout::MQALogitsSharedStorage< + {}, {}, {}, {}, {}, {}, {}, 3, {}, {} +>) == {}, "Incorrect MQA logits shared-memory size"); + static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sm100_paged_mqa_logits< - {}, {}, + auto ptr = reinterpret_cast(&sm100_mqa_logits< {}, {}, + {}, {}, {}, {}, {}, {}, {}, {}, + {}, + {}, {}, {}, {}, - {}, {}, {} + {} >); }}; -)", args.tokens_per_request, args.num_heads, - args.head_dim, args.page_kv, - args.is_mx_sf ? "true" : "false", - args.is_context_lens_2d, args.is_varlen ? "true" : "false", - args.num_q_stages, args.num_kv_stages, - args.split_kv, args.splits_per_chunk, - args.num_specialized_threads, args.num_math_threads, - args.qk_dtype == kPackedFP4 ? "cutlass::float_e2m1_t" : "cutlass::float_e4m3_t", - to_string(args.logits_dtype), - args.weights_dtype == torch::kBFloat16 ? "__nv_bfloat16" : "float"); - } +)", + num_heads, head_dim, is_mx_sf ? "true" : "false", + block_q, split_kv, num_q_stages, num_kv_stages, + qk_dtype == kPackedFP4 ? "cutlass::float_e2m1_t" : "cutlass::float_e4m3_t", + weights.scalar_type() == torch::kBFloat16 ? "__nv_bfloat16" : "float", + smem_size, + num_heads, head_dim, + is_mx_sf ? "true" : "false", + is_compressed_logits, clean_logits, + schedule_meta.has_value(), + block_q, split_kv, + num_q_stages, num_kv_stages, + num_sms, + num_specialized_threads, num_math_threads, + qk_dtype == kPackedFP4 ? "cutlass::float_e2m1_t" : "cutlass::float_e4m3_t", + to_string(logits_dtype), + weights.scalar_type() == torch::kBFloat16 ? "__nv_bfloat16" : "float")); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(num_specialized_threads + num_math_threads, 1, 1), + }, + num_q_tokens, num_kv_tokens, + stride_logits, + cu_seq_len_k_start.data_ptr(), cu_seq_len_k_end.data_ptr(), + schedule_meta.has_value() ? schedule_meta.value().data_ptr() : nullptr, + logits.data_ptr(), + tensor_map_q, tensor_map_sf_q, + tensor_map_kv, tensor_map_sf_kv, + tensor_map_weights + ); +} - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.num_q_tokens_total, - args.logits_stride, args.block_table_stride, - args.context_lens, args.logits, - args.block_table, args.indices, args.schedule_meta, - args.tensor_map_q, args.tensor_map_sf_q, - args.tensor_map_kv, args.tensor_map_sf_kv, - args.tensor_map_weights - )); - } -}; +// Paged variant: separate host/runtime path, shared device core +// Unified paged runtime for FP8 / MXFP4 / MXFP8; FP8 reuses the unused `sf_q` descriptor slot static void sm100_paged_mqa_logits(const torch::Tensor& q, const std::optional& sf_q, const torch::Tensor& kv_cache, @@ -411,8 +317,7 @@ static void sm100_paged_mqa_logits(const torch::Tensor& q, const int num_q_stages = 3; // Match contiguous-KV pipeline depth. const int num_kv_stages = is_fp4 ? 10 : 5; - // BLOCK_Q = 128 / num_heads; a Q-block holds up to BLOCK_Q request tokens - DG_HOST_ASSERT(128 % num_heads == 0); + DG_HOST_ASSERT(num_heads > 0 and num_heads <= 128 and num_heads % 4 == 0); const int block_q = 128 / num_heads; // MX SF formats consume `sf_q`; FP8 fills that descriptor slot with KV scales @@ -442,49 +347,72 @@ static void sm100_paged_mqa_logits(const torch::Tensor& q, } else { tensor_map_sf_q = tensor_map_sf_kv; // unused by FP8 } - const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, num_requests * tokens_per_request, - num_heads, block_q, - static_cast(weights.stride(0)), 0); - - const int smem_size = get_mqa_logits_smem_size(num_heads, head_dim, is_mx_sf, qk_dtype); - - const SM100PagedMQALogitsRuntime::Args args = { - .num_q_tokens_total = num_q_tokens_total, - .tokens_per_request = tokens_per_request, - .num_heads = num_heads, - .head_dim = head_dim, - .page_kv = page_kv, - .is_mx_sf = is_mx_sf, - .is_context_lens_2d = is_context_lens_2d, - .is_varlen = is_varlen, - .block_table_stride = block_table_stride, - .logits_stride = logits_stride, - .num_q_stages = num_q_stages, - .num_kv_stages = num_kv_stages, - .split_kv = split_kv, - .splits_per_chunk = splits_per_chunk, - .context_lens = context_lens.data_ptr(), - .logits = logits.data_ptr(), - .block_table = block_table.data_ptr(), - .indices = is_varlen ? indices.data_ptr() : nullptr, - .schedule_meta = schedule_meta.data_ptr(), - .tensor_map_q = tensor_map_q, - .tensor_map_sf_q = tensor_map_sf_q, - .tensor_map_kv = tensor_map_kv, - .tensor_map_sf_kv = tensor_map_sf_kv, - .tensor_map_weights = tensor_map_weights, - .qk_dtype = qk_dtype, - .logits_dtype = logits_dtype, - .weights_dtype = weights.scalar_type(), - .num_specialized_threads = num_specialized_threads, - .num_math_threads = num_math_threads, - .launch_args = LaunchArgs(num_sms, - num_specialized_threads + num_math_threads, - smem_size) - }; - const auto code = SM100PagedMQALogitsRuntime::generate(args); - const auto runtime = compiler->build("sm100_paged_mqa_logits", code); - SM100PagedMQALogitsRuntime::launch(runtime, args); + const uint32_t weights_stride = static_cast(weights.stride(0)); + const uint32_t num_weight_row_bytes = num_heads * static_cast(weights.element_size()); + const uint32_t num_weight_elements_per_row = math::align(num_weight_row_bytes, 16u) / static_cast(weights.element_size()); + const auto tensor_map_weights = make_tma_2d_desc( + weights, num_heads, num_requests * tokens_per_request, + static_cast(num_weight_elements_per_row), block_q, + static_cast(weights_stride), 0); + + const int smem_size = get_mqa_logits_smem_size( + num_heads, head_dim, is_mx_sf, qk_dtype, weights.scalar_type(), + block_q, split_kv, num_q_stages, num_kv_stages, 3); + DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity); + + // Compile + const auto kernel = jit->compile("sm100_paged_mqa_logits", std::format(R"( +#include + +using namespace deep_gemm; + +static_assert(sizeof(layout::MQALogitsSharedStorage< + {}, {}, {}, {}, {}, {}, {}, 3, {}, {} +>) == {}, "Incorrect MQA logits shared-memory size"); + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_paged_mqa_logits< + {}, {}, + {}, {}, + {}, {}, {}, + {}, {}, + {}, {}, + {}, {}, + {}, {}, {} + >); +}}; +)", + num_heads, head_dim, is_mx_sf ? "true" : "false", + block_q, split_kv, num_q_stages, num_kv_stages, + qk_dtype == kPackedFP4 ? "cutlass::float_e2m1_t" : "cutlass::float_e4m3_t", + weights.scalar_type() == torch::kBFloat16 ? "__nv_bfloat16" : "float", + smem_size, + tokens_per_request, num_heads, + head_dim, page_kv, + is_mx_sf ? "true" : "false", + is_context_lens_2d, is_varlen ? "true" : "false", + num_q_stages, num_kv_stages, + split_kv, splits_per_chunk, + num_specialized_threads, num_math_threads, + qk_dtype == kPackedFP4 ? "cutlass::float_e2m1_t" : "cutlass::float_e4m3_t", + to_string(logits_dtype), + weights.scalar_type() == torch::kBFloat16 ? "__nv_bfloat16" : "float")); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(num_specialized_threads + num_math_threads, 1, 1), + }, + num_q_tokens_total, + logits_stride, block_table_stride, + context_lens.data_ptr(), logits.data_ptr(), + block_table.data_ptr(), is_varlen ? indices.data_ptr() : nullptr, schedule_meta.data_ptr(), + tensor_map_q, tensor_map_sf_q, + tensor_map_kv, tensor_map_sf_kv, + tensor_map_weights + ); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm100_sparse_mqa_logits.hpp b/csrc/jit_kernels/impls/sm100_sparse_mqa_logits.hpp new file mode 100644 index 0000000000..c438547d72 --- /dev/null +++ b/csrc/jit_kernels/impls/sm100_sparse_mqa_logits.hpp @@ -0,0 +1,202 @@ +#pragma once + +#include + +#include + +#include "../../runtime/runtime.hpp" +#include "../heuristics/sm100.hpp" +#include "runtime_utils.hpp" + +namespace deep_gemm { + +using namespace layout::sparse_mqa_logits; + +static int get_sparse_mqa_split_kv(const at::ScalarType qk_dtype) { + DG_HOST_ASSERT(qk_dtype == kPackedFP4 or qk_dtype == torch::kFloat8_e4m3fn); + return qk_dtype == kPackedFP4 ? 640 : 512; +} + +static int64_t get_num_metadata_bytes(const int num_q_tokens, const int num_max_sparse_blocks, + const int sparse_block_kv, const int split_kv, + const bool is_paged, const int num_sms) { + DG_HOST_ASSERT(num_max_sparse_blocks > 0 and num_max_sparse_blocks % 4 == 0 and num_max_sparse_blocks <= 4096); + DG_HOST_ASSERT(sparse_block_kv == 8 or sparse_block_kv == 16); + const int64_t num_kv_blocks_per_split = split_kv / sparse_block_kv; + const int64_t num_kv_split_bytes = sizeof(KVSplitHeader) + num_kv_blocks_per_split * sizeof(KVBlockInfo); + const int64_t num_max_kv_splits = is_paged ? + static_cast(num_q_tokens) * ceil_div(num_max_sparse_blocks, num_kv_blocks_per_split) : + ceil_div(num_q_tokens, kBlockQ) * + ceil_div(kBlockQ * static_cast(num_max_sparse_blocks), num_kv_blocks_per_split); + DG_HOST_ASSERT(num_max_kv_splits <= std::numeric_limits::max()); + const int64_t num_max_schedule_entries = align(num_max_kv_splits, num_sms); + DG_HOST_ASSERT(num_max_schedule_entries <= std::numeric_limits::max()); + return sizeof(MetadataHeader) + num_max_kv_splits * num_kv_split_bytes + + num_max_schedule_entries * static_cast(sizeof(ScheduleEntry)); +} + +static void launch_sm100_sparse_mqa_logits(const bool is_paged, const bool use_unaligned_ks, + const int sparse_block_kv, + const torch::Tensor& q, const torch::Tensor& sf_q, + const torch::Tensor& kv, const torch::Tensor& sf_kv, + const torch::Tensor& weights, const torch::Tensor& metadata, + const torch::Tensor& logits) { + constexpr int kNumQStages = 2; + constexpr int kNumTmemStages = 5; + const bool is_fp4 = q.scalar_type() == kPackedFP4; + const int split_kv = get_sparse_mqa_split_kv(q.scalar_type()); + const int num_math_warpgroups = split_kv / 128; + const int num_kv_stages = is_fp4 ? 5 : 3; + const auto qk_dtype_name = is_fp4 ? "cutlass::float_e2m1_t" : "cutlass::float_e4m3_t"; + DG_HOST_ASSERT(not is_paged or not use_unaligned_ks); + DG_HOST_ASSERT(sparse_block_kv == 8 or sparse_block_kv == 16); + DG_HOST_ASSERT(metadata.dim() == 1 and metadata.scalar_type() == torch::kUInt8 and metadata.is_contiguous() and + metadata.numel() >= static_cast(sizeof(MetadataHeader))); + const int num_q_tokens = static_cast(q.size(0)); + const int swizzle_mode = is_fp4 ? kHeadDim / 2 : kHeadDim; + const auto tensor_map_q = make_tma_2d_desc(q, kHeadDim, num_q_tokens * kNumHeads, kHeadDim, kBlockQ * kNumHeads, + static_cast(q.stride(is_paged ? 2 : 1)), swizzle_mode, + 0, false, not is_fp4); + const auto tensor_map_sf_q = make_tma_2d_desc(sf_q, kNumHeads, num_q_tokens, kNumHeads, kBlockQ, + static_cast(sf_q.stride(is_paged ? 1 : 0)), 0); + const auto tensor_map_weights = make_tma_2d_desc(weights, kNumHeads, num_q_tokens, kNumHeads, kBlockQ, + static_cast(weights.stride(0)), 0); + CUtensorMap tensor_map_kv{}; + CUtensorMap tensor_map_sf_kv{}; + if (not is_paged) { + const int num_kv_tokens = static_cast(kv.size(0)); + tensor_map_kv = make_tma_2d_desc(kv, kHeadDim, num_kv_tokens, kHeadDim, kNumKVTokensPerTMA, + static_cast(kv.stride(0)), swizzle_mode, 0, false, not is_fp4); + tensor_map_sf_kv = make_tma_2d_desc(sf_kv, get_tma_aligned_size(num_kv_tokens, static_cast(sf_kv.element_size())), 1, + kNumKVTokensPerTMA, 1, 0, 0); + } + const int num_smem_bytes = is_fp4 ? + (sparse_block_kv == 8 ? + static_cast(sizeof(SharedStorage)) : + static_cast(sizeof(SharedStorage))) : + (sparse_block_kv == 8 ? + static_cast(sizeof(SharedStorage)) : + static_cast(sizeof(SharedStorage))); + DG_HOST_ASSERT(num_smem_bytes <= SM100ArchSpec::smem_capacity); + + const int num_sms = runtime->get_num_sms(); + const auto instantiate = is_paged ? std::format(R"( + auto ptr = reinterpret_cast(&sm100_paged_sparse_mqa_logits< + {}, {}, {}, {}, {}, {}, {}, {}, {} + >); +)", static_cast(kv.size(1)), sparse_block_kv, kNumQStages, num_kv_stages, kNumTmemStages, + num_math_warpgroups, num_sms, kBlockQ, is_fp4) : std::format(R"( + auto ptr = reinterpret_cast(&sm100_sparse_mqa_logits< + {}, {}, {}, {}, {}, {}, {}, {}, {} + >); +)", sparse_block_kv, kNumQStages, num_kv_stages, kNumTmemStages, num_math_warpgroups, num_sms, kBlockQ, + use_unaligned_ks, is_fp4); + const auto kernel_name = is_paged ? "sm100_paged_sparse_mqa_logits" : "sm100_sparse_mqa_logits"; + const auto kernel = jit->compile(kernel_name, std::format(R"( +#include + +using namespace deep_gemm; +using namespace deep_gemm::layout::sparse_mqa_logits; + +static_assert(sizeof(SharedStorage< + {}, {}, {}, {}, {}, {}, {} +>) == {}, "Incorrect sparse MQA logits shared-memory size"); + +static void __instantiate_kernel() {{ +{} +}} +)", kBlockQ, sparse_block_kv, split_kv, kNumQStages, num_kv_stages, kNumTmemStages, + qk_dtype_name, num_smem_bytes, instantiate)); + + if (is_paged) { + jit->launch( + kernel, { + .num_smem_bytes = num_smem_bytes, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(get_num_threads(num_math_warpgroups), 1, 1), + }, + static_cast(logits.stride(0)), static_cast(kv.stride(0)), logits.data_ptr(), + kv.data_ptr(), metadata.data_ptr(), tensor_map_q, tensor_map_sf_q, tensor_map_weights + ); + } else { + jit->launch( + kernel, { + .num_smem_bytes = num_smem_bytes, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(get_num_threads(num_math_warpgroups), 1, 1), + }, + static_cast(logits.stride(0)), logits.data_ptr(), kv.data_ptr(), sf_kv.data_ptr(), + metadata.data_ptr(), tensor_map_q, tensor_map_sf_q, tensor_map_weights, + tensor_map_kv, tensor_map_sf_kv + ); + } +} + +static void launch_sm100_sparse_mqa_logits_metadata(const bool is_paged, + const bool use_unaligned_ks, + const int page_kv, + const int num_kv_tokens, + const int block_table_stride, + const torch::Tensor& sparse_kv_block_indices, + const torch::Tensor& metadata, + const torch::Tensor& workspace, + const int split_kv, + const int sparse_block_kv, + const int* cu_seq_len_k_start, + const int* cu_seq_len_k_end, + const int* context_lens, + const int* block_table, + const int* indices) { + constexpr int kNumMetadataThreads = 256; + constexpr int kNumKVSplitsPerEntry = 8; + DG_HOST_ASSERT(not is_paged or not use_unaligned_ks); + const int num_q_tokens = static_cast(sparse_kv_block_indices.size(0)); + const int num_max_sparse_blocks = static_cast(sparse_kv_block_indices.size(1)); + const int num_sms = runtime->get_num_sms(); + const int num_ctas = std::min(is_paged ? num_q_tokens : ceil_div(num_q_tokens, kBlockQ), num_sms * 4); + const int num_max_merged_kv_blocks = kBlockQ * num_max_sparse_blocks; + const int num_max_kv_splits = ceil_div(num_max_merged_kv_blocks, split_kv / sparse_block_kv); + int num_smem_bytes = (kBlockQ * num_max_sparse_blocks + num_max_merged_kv_blocks + num_max_kv_splits) * + static_cast(sizeof(uint32_t)); + num_smem_bytes += (kNumMetadataThreads / 32) * static_cast(sizeof(uint32_t)); + num_smem_bytes += (1 + kNumMetadataThreads / 32 * (kNumKVSplitsPerEntry + 1)) * static_cast(sizeof(uint32_t)); + num_smem_bytes = align(num_smem_bytes, 16); + const int num_q_block_bytes = align((3 + kBlockQ) * static_cast(sizeof(uint32_t)), 16); + num_smem_bytes = align(num_smem_bytes + num_q_block_bytes + static_cast(sizeof(uint32_t)), 128); + DG_HOST_ASSERT(num_ctas > 0 and num_smem_bytes <= SM100ArchSpec::smem_capacity); + + const auto kernel = jit->compile("sm100_sparse_mqa_logits_metadata", std::format(R"( +#include + +using namespace deep_gemm; +using namespace deep_gemm::sched::sparse_mqa_logits; + +static void __instantiate_kernel() {{ + static_assert((sizeof(sched::sparse_mqa_logits::SharedStorage< + {}, {}, {} / {}, {}, {} + >) + 127) / 128 * 128 == {}, "Incorrect sparse MQA metadata shared-memory size"); + auto ptr = reinterpret_cast(&sm100_sparse_mqa_logits_metadata< + {}, {}, {}, {}, {}, {}, {}, {}, {}, {} + >); +}}; +)", kBlockQ, num_max_sparse_blocks, split_kv, sparse_block_kv, + kNumKVSplitsPerEntry, kNumMetadataThreads, num_smem_bytes, + is_paged ? "true" : "false", use_unaligned_ks ? "true" : "false", + kBlockQ, split_kv, sparse_block_kv, + num_max_sparse_blocks, page_kv, kNumKVSplitsPerEntry, num_sms, kNumMetadataThreads)); + + jit->launch( + kernel, { + .num_smem_bytes = num_smem_bytes, + .grid_dim = dim3(num_ctas, 1, 1), + .block_dim = dim3(kNumMetadataThreads, 1, 1), + }, + num_q_tokens, num_kv_tokens, + cu_seq_len_k_start, cu_seq_len_k_end, + context_lens, block_table, block_table_stride, + indices, + sparse_kv_block_indices.data_ptr(), metadata.data_ptr(), workspace.data_ptr() + ); +} + +} // namespace deep_gemm 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..004f73070b 100644 --- a/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp @@ -1,68 +1,17 @@ #pragma once +#include +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/jit.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../../utils/math.hpp" #include "../heuristics/sm100.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM100BF16HCPrenormGemmRuntime final: public LaunchRuntime { -public: - struct Args { - int m, n, k; - int block_m, block_n, block_k; - int num_splits; - int swizzle_cd_mode; - int num_stages; - int num_mma_threads, num_cast_and_reduce_threads; - - LaunchArgs launch_args; - - CUtensorMap tensor_map_a; - CUtensorMap tensor_map_b; - CUtensorMap tensor_map_d; - float* sqr_sum; - }; - - 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(&sm100_tf32_hc_prenorm_gemm_impl< - {}, {}, - {}, {}, {}, - {}, - {}, - {}, - {}, {} - >); -}}; -)", - args.n, args.k, - args.block_m, args.block_n, args.block_k, - args.num_splits, - args.swizzle_cd_mode, - args.num_stages, - args.num_mma_threads, args.num_cast_and_reduce_threads); - } - - 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.m, args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.sqr_sum)); - } -}; - static void sm100_tf32_hc_prenorm_gemm(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, @@ -118,7 +67,7 @@ static void sm100_tf32_hc_prenorm_gemm(const torch::Tensor& a, DG_HOST_ASSERT(num_stages > 0); // Print configs - if (get_env("DG_JIT_DEBUG", 0)) { + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { printf("M: %d, N: %d, K: %d -> " "block M: %d, block N: %d, block K: %d, split K: %d" "stages: %d, shared memory: %d, swizzle CD: %d\n", @@ -126,24 +75,39 @@ static void sm100_tf32_hc_prenorm_gemm(const torch::Tensor& a, num_stages, smem_size, swizzle_cd_mode); } + // Compile + const auto kernel = jit->compile("sm100_tf32_hc_prenorm_gemm", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_tf32_hc_prenorm_gemm_impl< + {}, {}, + {}, {}, {}, + {}, + {}, + {}, + {}, {} + >); +}}; +)", + n, k, + block_m, block_n, block_k, + num_splits, + swizzle_cd_mode, + num_stages, + num_mma_threads, num_cast_and_reduce_threads)); + // Launch - const SM100BF16HCPrenormGemmRuntime::Args& args = { - .m = m, .n = n, .k = k, - .block_m = block_m, .block_n = block_n, .block_k = block_k, - .num_splits = num_splits, - .swizzle_cd_mode = swizzle_cd_mode, - .num_stages = num_stages, - .num_mma_threads = num_mma_threads, - .num_cast_and_reduce_threads = num_cast_and_reduce_threads, - .launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_mma_threads + num_cast_and_reduce_threads, smem_size), - .tensor_map_a = tensor_map_a, - .tensor_map_b = tensor_map_b, - .tensor_map_d = tensor_map_d, - .sqr_sum = sqr_sum.data_ptr() - }; - const auto code = SM100BF16HCPrenormGemmRuntime::generate(args); - const auto runtime = compiler->build("sm100_tf32_hc_prenorm_gemm", code); - SM100BF16HCPrenormGemmRuntime::launch(runtime, args); + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_splits * ceil_div(m, block_m), 1, 1), + .block_dim = dim3(num_mma_threads + num_cast_and_reduce_threads, 1, 1), + }, + m, tensor_map_a, tensor_map_b, tensor_map_d, sqr_sum.data_ptr() + ); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp index 24edd46562..699db4191e 100644 --- a/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_bf16_gemm.hpp @@ -1,22 +1,21 @@ #pragma once +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../heuristics/sm90.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM90BF16GemmRuntime final: public LaunchRuntime { +class SM90BF16GemmRuntime final { public: struct Args { GemmDesc gemm_desc; GemmConfig gemm_config; - LaunchArgs launch_args; + deep_jit::cuda::LaunchOptions options; void *grouped_layout; CUtensorMap tensor_map_a; @@ -24,8 +23,8 @@ class SM90BF16GemmRuntime final: public LaunchRuntime { CUtensorMap tensor_map_cd; }; - static std::string generate_impl(const Args& args) { - return fmt::format(R"( + static void compile_and_launch(const std::string& tag, const Args& args) { + const auto kernel = jit->compile(tag, std::format(R"( #include using namespace deep_gemm; @@ -62,16 +61,16 @@ static void __instantiate_kernel() {{ 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), args.gemm_desc.with_accumulation, - to_string(args.gemm_desc.cd_dtype)); - } + to_string(args.gemm_desc.cd_dtype))); - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - // TODO: optimize `args` copy - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + // Launch + jit->launch( + kernel, args.options, args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, args.tensor_map_a, args.tensor_map_b, - args.tensor_map_cd)); + args.tensor_map_cd + ); } }; @@ -90,8 +89,9 @@ static void sm90_bf16_gemm(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -112,21 +112,21 @@ static void sm90_bf16_gemm(const torch::Tensor& a, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); - // Launch - const SM90BF16GemmRuntime::Args& args = { + // Compile and launch + SM90BF16GemmRuntime::compile_and_launch("sm90_bf16_gemm", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM90BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm90_bf16_gemm", code); - SM90BF16GemmRuntime::launch(runtime, args); + }); } static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, @@ -157,8 +157,9 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, .expected_m = expected_m_for_psum_layout.value_or(m), .expected_n = n, .expected_k = k, .expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1 @@ -182,21 +183,21 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); - // Launch - const SM90BF16GemmRuntime::Args& args = { + // Compile and launch + SM90BF16GemmRuntime::compile_and_launch("sm90_m_grouped_bf16_gemm_contiguous", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = m_indices.data_ptr(), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM90BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_bf16_gemm_contiguous", code); - SM90BF16GemmRuntime::launch(runtime, args); + }); } static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a, @@ -219,8 +220,9 @@ static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, .expected_m = expected_m, .expected_n = 0, .expected_k = 0, .expected_num_groups = num_groups }; const auto config = get_best_config(desc); @@ -242,21 +244,21 @@ static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a, static_cast(d.stride(-2)), num_groups, config.storage_config.swizzle_cd_mode); - // Launch - const SM90BF16GemmRuntime::Args& args = { + // Compile and launch + SM90BF16GemmRuntime::compile_and_launch("sm90_bf16_m_grouped_gemm_masked", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = masked_m.data_ptr(), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM90BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm90_bf16_m_grouped_gemm_masked", code); - SM90BF16GemmRuntime::launch(runtime, args); + }); } static void sm90_bf16_k_grouped_gemm(const torch::Tensor& a, @@ -286,8 +288,9 @@ static void sm90_bf16_k_grouped_gemm(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, .expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups }; const auto config = get_best_config(desc); @@ -309,21 +312,21 @@ static void sm90_bf16_k_grouped_gemm(const torch::Tensor& a, static_cast(d.stride(1)), num_groups, config.storage_config.swizzle_cd_mode); - // Launch kernel - const SM90BF16GemmRuntime::Args& args = { + // Compile and launch + SM90BF16GemmRuntime::compile_and_launch("sm90_bf16_k_grouped_gemm", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = grouped_layout.data_ptr(), .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM90BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm90_bf16_k_grouped_gemm", code); - SM90BF16GemmRuntime::launch(runtime, args); + }); } static void sm90_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, @@ -339,8 +342,9 @@ static void sm90_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, .cd_dtype = tensor_d.scalar_type(), .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -361,21 +365,21 @@ static void sm90_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, tensor_d.stride(0), tensor_d.stride(1), config.storage_config.swizzle_cd_mode); - // Launch - const SM90BF16GemmRuntime::Args& args = { + // Compile and launch + SM90BF16GemmRuntime::compile_and_launch("sm90_bf16_bhr_hdr_bhd", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM90BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm90_bf16_bhr_hdr_bhd", code); - SM90BF16GemmRuntime::launch(runtime, args); + }); } static void sm90_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, @@ -391,8 +395,9 @@ static void sm90_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, .cd_dtype = tensor_d.scalar_type(), .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::MN, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -412,21 +417,22 @@ static void sm90_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, store_block_n, store_block_m, 1, tensor_d.stride(0), tensor_d.stride(1), config.storage_config.swizzle_cd_mode); - // Launch - const SM90BF16GemmRuntime::Args& args = { + + // Compile and launch + SM90BF16GemmRuntime::compile_and_launch("sm90_bf16_bhd_hdr_bhr", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .grouped_layout = nullptr, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM90BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm90_bf16_bhd_hdr_bhr", code); - SM90BF16GemmRuntime::launch(runtime, args); + }); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp index 473677b70c..2bda372666 100644 --- a/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp @@ -1,64 +1,17 @@ #pragma once +#include +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../../utils/math.hpp" #include "../heuristics/sm90.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM90BmkBnkMnRuntime final: public LaunchRuntime { -public: - struct Args { - int s, m, n, k; - int block_m, block_n, block_k; - int split_factor; - int num_stages; - int num_tma_threads, num_math_threads; - - LaunchArgs launch_args; - - CUtensorMap tensor_map_a; - CUtensorMap tensor_map_b; - float* d; - }; - - 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_bmn_bnk_mn_gemm_impl< - {}, {}, {}, - {}, {}, {}, - {}, - {}, - {}, {} - >); -}}; -)", - args.m, args.n, args.k, - args.block_m, args.block_n, args.block_k, - args.split_factor, - args.num_stages, - args.num_tma_threads, args.num_math_threads); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.s, args.tensor_map_a, args.tensor_map_b, args.d)); - } -}; - - static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a, const torch::Tensor &b, const torch::Tensor &d, @@ -76,7 +29,7 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a, DG_HOST_ASSERT(swizzle_ab_mode == 128); // Get best config - const int num_sms = device_runtime->get_num_sms(); + const int num_sms = runtime->get_num_sms(); const int num_mn_blocks = ceil_div(m, block_m) * ceil_div(n, block_n); const int num_sk_blocks = s * (k / block_k); const int split_factor = ceil_div(num_sk_blocks, std::max(num_sms / num_mn_blocks, 1)); @@ -100,7 +53,7 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a, DG_HOST_ASSERT(num_stages > 0); // Print configs - if (get_env("DG_JIT_DEBUG", 0)) { + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { printf("S: %d, M: %d, N: %d, K: %d -> " "block M: %d, block N: %d, block K: %d, split-K factor: %d" "stages: %d, shared memory: %d, swizzle AB: %d\n", @@ -111,21 +64,37 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a, const auto tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode); const auto tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode); - const SM90BmkBnkMnRuntime::Args& args = { - .s = s, .m = m, .n = n, .k = k, - .block_m = block_m, .block_n = block_n, .block_k = block_k, - .split_factor = split_factor, - .num_stages = num_stages, - .num_tma_threads = num_tma_threads, - .num_math_threads = num_math_threads, - .launch_args = LaunchArgs(num_mn_blocks * ceil_div(num_sk_blocks, split_factor), num_tma_threads + num_math_threads, smem_size), - .tensor_map_a = tensor_map_a, - .tensor_map_b = tensor_map_b, - .d = d.data_ptr() - }; - const auto code = SM90BmkBnkMnRuntime::generate(args); - const auto runtime = compiler->build("sm90_bmn_bnk_mn_gemm", code); - SM90BmkBnkMnRuntime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("sm90_bmn_bnk_mn_gemm", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_bmn_bnk_mn_gemm_impl< + {}, {}, {}, + {}, {}, {}, + {}, + {}, + {}, {} + >); +}}; +)", + m, n, k, + block_m, block_n, block_k, + split_factor, + num_stages, + num_tma_threads, num_math_threads)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_mn_blocks * ceil_div(num_sk_blocks, split_factor), 1, 1), + .block_dim = dim3(num_tma_threads + num_math_threads, 1, 1), + }, + s, tensor_map_a, tensor_map_b, d.data_ptr() + ); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp index 120b91faf8..29f516c6d6 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp @@ -1,23 +1,21 @@ #pragma once +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../heuristics/sm90.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM90FP8Gemm1D1DRuntime final: public LaunchRuntime { +class SM90FP8Gemm1D1DRuntime final { public: struct Args { GemmDesc gemm_desc; GemmConfig gemm_config; - LaunchArgs launch_args; + deep_jit::cuda::LaunchOptions options; void *gmem_a_ptr; void *gmem_b_ptr; @@ -30,8 +28,8 @@ class SM90FP8Gemm1D1DRuntime final: public LaunchRuntime CUtensorMap tensor_map_cd; }; - static std::string generate_impl(const Args& args) { - return fmt::format(R"( + static void compile_and_launch(const std::string& tag, const Args& args) { + const auto kernel = jit->compile(tag, std::format(R"( #include using namespace deep_gemm; @@ -60,18 +58,19 @@ 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), - to_string(args.gemm_desc.cd_dtype)); - } + to_string(args.gemm_desc.cd_dtype))); - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + // Launch + jit->launch( + kernel, args.options, args.gmem_a_ptr, args.gmem_b_ptr, args.grouped_layout, args.tensor_map_buffer, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k, args.tensor_map_a_base, args.tensor_map_b_base, args.tensor_map_sfa, args.tensor_map_sfb, - args.tensor_map_cd)); + args.tensor_map_cd + ); } }; @@ -93,8 +92,9 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -120,13 +120,16 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, static_cast(d.stride(-2)), 1, 0); - // Launch - const SM90FP8Gemm1D1DRuntime::Args& args = { + // Compile and launch + SM90FP8Gemm1D1DRuntime::compile_and_launch("sm90_fp8_gemm_1d1d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .gmem_a_ptr = nullptr, .gmem_b_ptr = nullptr, .grouped_layout = nullptr, @@ -136,11 +139,7 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM90FP8Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm90_fp8_gemm_1d1d", code); - - SM90FP8Gemm1D1DRuntime::launch(runtime, args); + }); } static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -175,8 +174,9 @@ static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Te .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, .expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups }; const auto config = get_best_config(desc); @@ -203,13 +203,16 @@ static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Te static_cast(d.stride(-2)), num_groups, config.storage_config.swizzle_cd_mode); - // Launch - const SM90FP8Gemm1D1DRuntime::Args& args = { + // Compile and launch + SM90FP8Gemm1D1DRuntime::compile_and_launch("sm90_fp8_gemm_1d1d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .gmem_a_ptr = a.data_ptr(), .gmem_b_ptr = b.data_ptr(), .grouped_layout = grouped_layout.data_ptr(), @@ -219,11 +222,7 @@ static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Te .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM90FP8Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm90_fp8_gemm_1d1d", code); - - SM90FP8Gemm1D1DRuntime::launch(runtime, args); + }); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp index 892edaee7d..89d0d2b2fe 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -1,12 +1,10 @@ #pragma once +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../heuristics/sm90.hpp" #include "epilogue.hpp" @@ -14,12 +12,12 @@ namespace deep_gemm { -class SM90FP8Gemm1D2DRuntime final: public LaunchRuntime { +class SM90FP8Gemm1D2DRuntime final { public: struct Args { GemmDesc gemm_desc; GemmConfig gemm_config; - LaunchArgs launch_args; + deep_jit::cuda::LaunchOptions options; // TODO: move this into `gemm_desc` const std::optional& epilogue_type; @@ -31,8 +29,8 @@ class SM90FP8Gemm1D2DRuntime final: public LaunchRuntime CUtensorMap tensor_map_sfa; }; - static std::string generate_impl(const Args& args) { - return fmt::format(R"( + static void compile_and_launch(const std::string& tag, const Args& args) { + const auto kernel = jit->compile(tag, std::format(R"( #include using namespace deep_gemm; @@ -65,16 +63,16 @@ static void __instantiate_kernel() {{ 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), to_string(args.gemm_desc.cd_dtype), - get_default_epilogue_type(args.epilogue_type)); - } + get_default_epilogue_type(args.epilogue_type))); - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - // TODO: optimize `args` copy - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, + // Launch + jit->launch( + kernel, args.options, args.sfb, args.grouped_layout, 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)); + args.tensor_map_d, args.tensor_map_sfa + ); } }; @@ -97,8 +95,9 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -123,13 +122,16 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, config.layout.block_m, config.layout.block_k, 1, 0); - // Launch - const SM90FP8Gemm1D2DRuntime::Args& args = { + // Compile and launch + SM90FP8Gemm1D2DRuntime::compile_and_launch("sm90_fp8_gemm_1d2d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .epilogue_type = epilogue_type, .major_sfb = major_sfb, .sfb = sfb.data_ptr(), @@ -138,10 +140,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, .tensor_map_sfa = tensor_map_sfa, - }; - const auto code = SM90FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_fp8_gemm_1d2d", code); - SM90FP8Gemm1D2DRuntime::launch(runtime, args); + }); } static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -171,8 +170,9 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, .expected_m = expected_m_for_psum_layout.value_or(m), .expected_n = n, .expected_k = k, .expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1 @@ -200,13 +200,16 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, config.layout.block_m, config.layout.block_k, 1, 0); - // Launch - const SM90FP8Gemm1D2DRuntime::Args& args = { + // Compile and launch + SM90FP8Gemm1D2DRuntime::compile_and_launch("sm90_m_grouped_fp8_gemm_contiguous_1d2d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .epilogue_type = std::nullopt, .major_sfb = major_sfb, .sfb = sfb.data_ptr(), @@ -215,10 +218,7 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, .tensor_map_sfa = tensor_map_sfa, - }; - const auto code = SM90FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_m_grouped_fp8_gemm_contiguous_1d2d", code); - SM90FP8Gemm1D2DRuntime::launch(runtime, args); + }); } static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -240,8 +240,9 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = false, - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims, + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims, .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups }; const auto config = get_best_config(desc); @@ -267,13 +268,16 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, config.layout.block_m, config.layout.block_k, num_groups, 0); - // Launch - const SM90FP8Gemm1D2DRuntime::Args& args = { + // Compile and launch + SM90FP8Gemm1D2DRuntime::compile_and_launch("sm90_fp8_m_grouped_gemm_masked_1d2d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .epilogue_type = std::nullopt, .major_sfb = major_sfb, .sfb = sfb.data_ptr(), @@ -282,10 +286,7 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, .tensor_map_sfa = tensor_map_sfa, - }; - const auto code = SM90FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_fp8_m_grouped_gemm_masked_1d2d", code); - SM90FP8Gemm1D2DRuntime::launch(runtime, args); + }); } static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, @@ -306,8 +307,9 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, .cd_dtype = d.scalar_type(), .major_a = major_a, .major_b = major_b, .with_accumulation = c.has_value(), - .num_sms = device_runtime->get_num_sms(), - .tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), + .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -338,13 +340,16 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k, config.layout.block_m, config.layout.block_k, batch_size, 0); - // Launch - const SM90FP8Gemm1D2DRuntime::Args& args = { + // Compile and launch + SM90FP8Gemm1D2DRuntime::compile_and_launch("sm90_fp8_gemm_1d2d", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - config.layout.get_cluster_size()), + .options = { + .num_smem_bytes = config.pipeline_config.smem_size, + .grid_dim = dim3(config.launch_config.num_sms, 1, 1), + .block_dim = dim3(config.launch_config.num_threads, 1, 1), + .cluster_dim = dim3(config.layout.get_cluster_size(), 1, 1), + }, .epilogue_type = std::nullopt, .major_sfb = major_sfb, .sfb = sfb.data_ptr(), @@ -353,10 +358,7 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, .tensor_map_b = tensor_map_b, .tensor_map_d = tensor_map_d, .tensor_map_sfa = tensor_map_sfa, - }; - const auto code = SM90FP8Gemm1D2DRuntime::generate(args); - const auto runtime = compiler->build("sm90_fp8_gemm_1d2d", code); - SM90FP8Gemm1D2DRuntime::launch(runtime, args); + }); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm90_fp8_mqa_logits.hpp b/csrc/jit_kernels/impls/sm90_fp8_mqa_logits.hpp index a2b9e0fb0e..c2faf8a4c7 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mqa_logits.hpp @@ -1,85 +1,13 @@ #pragma once -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include + +#include "../../runtime/runtime.hpp" #include "../heuristics/sm90.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM90FP8MQALogitsRuntime final: public LaunchRuntime { -public: - struct Args { - int seq_len; - int seq_len_kv; - int max_seqlen_k; - int stride_logits; - int num_heads, head_dim; - bool is_compressed_logits; - - int num_q_stages; - int num_kv_stages; - int block_q; - int block_kv; - - int* cu_seq_len_k_start; - int* cu_seq_len_k_end; - void* logits; - - CUtensorMap tensor_map_q; - CUtensorMap tensor_map_kv; - CUtensorMap tensor_map_kv_scales; - CUtensorMap tensor_map_weights; - at::ScalarType logits_dtype; - - int num_specialized_threads; - int num_math_threads; - - LaunchArgs launch_args; - }; - - static std::string generate_impl(const Args& args) { - DG_HOST_ASSERT(128 % args.num_heads == 0); - - return fmt::format(R"( -#include - -using namespace deep_gemm; - -static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sm90_fp8_mqa_logits< - {}, {}, - {}, - {}, {}, - {}, {}, - {}, - {}, {}, - {} - >); -}}; -)", - args.num_heads, args.head_dim, - args.is_compressed_logits, - args.block_q, args.block_kv, - args.num_q_stages, args.num_kv_stages, - args.launch_args.grid_dim.first, - args.num_specialized_threads, args.num_math_threads, - to_string(args.logits_dtype)); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.seq_len, args.seq_len_kv, - args.max_seqlen_k, args.stride_logits, - args.cu_seq_len_k_start, args.cu_seq_len_k_end, - args.logits, - args.tensor_map_q, args.tensor_map_kv, - args.tensor_map_kv_scales, args.tensor_map_weights - )); - } -}; - static void sm90_fp8_mqa_logits(const torch::Tensor& q, const torch::Tensor& kv, const torch::Tensor& kv_scales, const torch::Tensor& weights, @@ -90,14 +18,17 @@ static void sm90_fp8_mqa_logits(const torch::Tensor& q, const int& seq_len, const int& seq_len_kv, const int& max_seqlen_k, const int& stride_logits, const int& num_heads, const int& head_dim, - const int& block_q, const int& block_kv) { + const int& block_q, const int& block_kv, + const bool& clean_logits) { constexpr int num_specialized_threads = 128; constexpr int num_q_stages = 3, num_kv_stages = 3; constexpr int num_math_threads = 512; const bool is_compressed_logits = (max_seqlen_k > 0); + const int num_sms = runtime->get_num_sms(); + DG_HOST_ASSERT(not (clean_logits and is_compressed_logits)); - DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + DG_HOST_ASSERT(jit->device.get_arch_major() == 9); DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128); const auto tensor_map_q = make_tma_2d_desc(q, head_dim, seq_len * num_heads, head_dim, block_q * num_heads, head_dim, head_dim); @@ -123,79 +54,49 @@ static void sm90_fp8_mqa_logits(const torch::Tensor& q, smem_size += 4; DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity); - const SM90FP8MQALogitsRuntime::Args args = { - .seq_len = seq_len, - .seq_len_kv = seq_len_kv, - .max_seqlen_k = max_seqlen_k, - .stride_logits = stride_logits, - .num_heads = num_heads, .head_dim = head_dim, - .is_compressed_logits = is_compressed_logits, - .num_q_stages = num_q_stages, - .num_kv_stages = num_kv_stages, - .block_q = block_q, - .block_kv = block_kv, - .cu_seq_len_k_start = cu_seq_len_k_start.data_ptr(), - .cu_seq_len_k_end = cu_seq_len_k_end.data_ptr(), - .logits = logits.data_ptr(), - .tensor_map_q = tensor_map_q, - .tensor_map_kv = tensor_map_kv, - .tensor_map_kv_scales = tensor_map_kv_scales, - .tensor_map_weights = tensor_map_weights, - .logits_dtype = logits_dtype, - .num_specialized_threads = num_specialized_threads, - .num_math_threads = num_math_threads, - .launch_args = LaunchArgs(device_runtime->get_num_sms(), - num_specialized_threads + num_math_threads, - smem_size) - }; - const auto code = SM90FP8MQALogitsRuntime::generate(args); - const auto runtime = compiler->build("sm90_fp8_mqa_logits", code); - SM90FP8MQALogitsRuntime::launch(runtime, args); -} - -class SM90PagedMQALogitsMetadataRuntime final: public LaunchRuntime { -public: - struct Args { - int aligned_batch_size; - int split_kv; - int num_sms; - bool is_varlen; + DG_HOST_ASSERT(128 % num_heads == 0); - int batch_size; - int next_n; - bool is_context_lens_2d; - int* context_lens; - int* indices; - int* schedule_metadata; - - LaunchArgs launch_args; - }; - - static std::string generate_impl(const Args& args) { - return fmt::format(R"( -#include + // Compile + const auto kernel = jit->compile("sm90_fp8_mqa_logits", std::format(R"( +#include using namespace deep_gemm; static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sched::sm90_paged_mqa_logits_metadata< - {}, {}, {}, {} + auto ptr = reinterpret_cast(&sm90_fp8_mqa_logits< + {}, {}, + {}, {}, + {}, {}, + {}, {}, + {}, + {}, {}, + {} >); }}; -)", args.aligned_batch_size, args.split_kv, args.num_sms, args.is_varlen ? "true" : "false"); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.batch_size, - args.next_n, - args.is_context_lens_2d, - args.context_lens, - args.indices, - args.schedule_metadata - )); - } -}; +)", + num_heads, head_dim, + is_compressed_logits, clean_logits, + block_q, block_kv, + num_q_stages, num_kv_stages, + num_sms, + num_specialized_threads, num_math_threads, + to_string(logits_dtype))); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(num_specialized_threads + num_math_threads, 1, 1), + }, + seq_len, seq_len_kv, + max_seqlen_k, stride_logits, + cu_seq_len_k_start.data_ptr(), cu_seq_len_k_end.data_ptr(), + logits.data_ptr(), + tensor_map_q, tensor_map_kv, + tensor_map_kv_scales, tensor_map_weights + ); +} static void sm90_paged_mqa_logits_metadata(const torch::Tensor& context_lens, const torch::Tensor& schedule_metadata, @@ -212,99 +113,34 @@ static void sm90_paged_mqa_logits_metadata(const torch::Tensor& context_lens, const int smem_size = num_smem_ints * static_cast(sizeof(int)); DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity); - const SM90PagedMQALogitsMetadataRuntime::Args& args = { - .aligned_batch_size = aligned_batch_size, - .split_kv = split_kv, - .num_sms = num_sms, - .is_varlen = is_varlen, - .batch_size = batch_size, - .next_n = next_n, - .is_context_lens_2d = is_context_lens_2d, - .context_lens = context_lens.data_ptr(), - .indices = const_cast(indices_ptr), - .schedule_metadata = schedule_metadata.data_ptr(), - .launch_args = LaunchArgs(1, num_threads, smem_size) - }; - const auto code = SM90PagedMQALogitsMetadataRuntime::generate(args); - const auto runtime = compiler->build("sm90_paged_mqa_logits_metadata", code); - SM90PagedMQALogitsMetadataRuntime::launch(runtime, args); -} - -class SM90FP8PagedMQALogitsRuntime final: public LaunchRuntime { -public: - struct Args { - int batch_size; - int next_n; - int num_heads; - int head_dim; - int block_kv; - bool is_context_lens_2d; - bool is_varlen; - int block_table_stride; - int logits_stride; - - int num_q_stages; - int num_kv_stages; - int split_kv; - - int* context_lens; - void* logits; - int* block_table; - int* indices; - int* schedule_meta; - - CUtensorMap tensor_map_q; - CUtensorMap tensor_map_kv; - CUtensorMap tensor_map_kv_scales; - CUtensorMap tensor_map_weights; - at::ScalarType logits_dtype; - - int num_specialized_threads; - int num_math_threads; - - LaunchArgs launch_args; - }; - - static std::string generate_impl(const Args& args) { - DG_HOST_ASSERT(128 % args.num_heads == 0); - - return fmt::format(R"( -#include + // Compile + const auto kernel = jit->compile("sm90_paged_mqa_logits_metadata", std::format(R"( +#include using namespace deep_gemm; static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sm90_fp8_paged_mqa_logits< - {}, {}, - {}, {}, - {}, {}, - {}, {}, - {}, - {}, {}, - {} + auto ptr = reinterpret_cast(&sched::sm90_paged_mqa_logits_metadata< + {}, {}, {}, {} >); }}; -)", - args.next_n, args.num_heads, - args.head_dim, args.block_kv, - args.is_context_lens_2d, args.is_varlen ? "true" : "false", - args.num_q_stages, args.num_kv_stages, - args.split_kv, - args.num_specialized_threads, args.num_math_threads, - to_string(args.logits_dtype)); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.batch_size, - args.logits_stride, args.block_table_stride, - args.context_lens, args.logits, - args.block_table, args.indices, args.schedule_meta, - args.tensor_map_q, args.tensor_map_kv, - args.tensor_map_kv_scales, args.tensor_map_weights - )); - } -}; +)", aligned_batch_size, split_kv, num_sms, is_varlen ? "true" : "false")); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(1, 1, 1), + .block_dim = dim3(num_threads, 1, 1), + }, + batch_size, + next_n, + is_context_lens_2d, + context_lens.data_ptr(), + const_cast(indices_ptr), + schedule_metadata.data_ptr() + ); +} static void sm90_fp8_paged_mqa_logits(const torch::Tensor& q, const torch::Tensor& kv_cache, @@ -330,7 +166,7 @@ static void sm90_fp8_paged_mqa_logits(const torch::Tensor& q, const int num_math_warp_groups = split_kv / mma_m; const int num_math_threads = num_math_warp_groups * 128; constexpr int num_q_stages = 3, num_kv_stages = 3; - DG_HOST_ASSERT(device_runtime->get_arch_major() == 9); + DG_HOST_ASSERT(jit->device.get_arch_major() == 9); DG_HOST_ASSERT(split_kv % mma_m == 0 and logits_stride % split_kv == 0); const int next_n_atom = (is_varlen or next_n >= 2) ? 2 : 1; @@ -363,38 +199,48 @@ static void sm90_fp8_paged_mqa_logits(const torch::Tensor& q, DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity); DG_HOST_ASSERT(next_n == 1 or next_n == 2); - const SM90FP8PagedMQALogitsRuntime::Args args = { - .batch_size = batch_size, - .next_n = next_n, - .num_heads = num_heads, - .head_dim = head_dim, - .block_kv = block_kv, - .is_context_lens_2d = is_context_lens_2d, - .is_varlen = is_varlen, - .block_table_stride = block_table_stride, - .logits_stride = logits_stride, - .num_q_stages = num_q_stages, - .num_kv_stages = num_kv_stages, - .split_kv = split_kv, - .context_lens = context_lens.data_ptr(), - .logits = logits.data_ptr(), - .block_table = block_table.data_ptr(), - .indices = is_varlen ? indices.data_ptr() : nullptr, - .schedule_meta = schedule_meta.data_ptr(), - .tensor_map_q = tensor_map_q, - .tensor_map_kv = tensor_map_kv, - .tensor_map_kv_scales = tensor_map_kv_scales, - .tensor_map_weights = tensor_map_weights, - .logits_dtype = logits_dtype, - .num_specialized_threads = num_specialized_threads, - .num_math_threads = num_math_threads, - .launch_args = LaunchArgs(num_sms, - num_specialized_threads + num_math_threads, - smem_size) - }; - const auto code = SM90FP8PagedMQALogitsRuntime::generate(args); - const auto runtime = compiler->build("sm90_fp8_paged_mqa_logits", code); - SM90FP8PagedMQALogitsRuntime::launch(runtime, args); + DG_HOST_ASSERT(128 % num_heads == 0); + + // Compile + const auto kernel = jit->compile("sm90_fp8_paged_mqa_logits", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_fp8_paged_mqa_logits< + {}, {}, + {}, {}, + {}, {}, + {}, {}, + {}, + {}, {}, + {} + >); +}}; +)", + next_n, num_heads, + head_dim, block_kv, + is_context_lens_2d, is_varlen ? "true" : "false", + num_q_stages, num_kv_stages, + split_kv, + num_specialized_threads, num_math_threads, + to_string(logits_dtype))); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_sms, 1, 1), + .block_dim = dim3(num_specialized_threads + num_math_threads, 1, 1), + }, + batch_size, + logits_stride, block_table_stride, + context_lens.data_ptr(), logits.data_ptr(), + block_table.data_ptr(), is_varlen ? indices.data_ptr() : nullptr, schedule_meta.data_ptr(), + tensor_map_q, tensor_map_kv, + tensor_map_kv_scales, tensor_map_weights + ); } } // namespace deep_gemm 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..2dcee9a0e3 100644 --- a/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp @@ -1,68 +1,17 @@ #pragma once +#include +#include #include -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include "../../runtime/jit.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../../utils/math.hpp" #include "../heuristics/sm90.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM90BF16HCPrenormGemmRuntime final: public LaunchRuntime { -public: - struct Args { - int m, n, k; - int block_m, block_n, block_k; - int num_splits; - int swizzle_cd_mode; - int num_stages; - int num_math_threads, num_tma_threads; - - LaunchArgs launch_args; - - CUtensorMap tensor_map_a; - CUtensorMap tensor_map_b; - CUtensorMap tensor_map_d; - float* sqr_sum; - }; - - 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_tf32_hc_prenorm_gemm_impl< - {}, {}, - {}, {}, {}, - {}, - {}, - {}, - {}, {} - >); -}}; -)", - args.n, args.k, - args.block_m, args.block_n, args.block_k, - args.num_splits, - args.swizzle_cd_mode, - args.num_stages, - args.num_math_threads, args.num_tma_threads); - } - - 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.m, args.tensor_map_a, args.tensor_map_b, args.tensor_map_d, args.sqr_sum)); - } -}; - static void sm90_tf32_hc_prenorm_gemm(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, @@ -119,7 +68,7 @@ static void sm90_tf32_hc_prenorm_gemm(const torch::Tensor& a, DG_HOST_ASSERT(num_stages > 0); // Print configs - if (get_env("DG_JIT_DEBUG", 0)) { + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { printf("M: %d, N: %d, K: %d -> " "block M: %d, block N: %d, block K: %d, split K: %d" "stages: %d, shared memory: %d, swizzle CD: %d\n", @@ -129,24 +78,39 @@ static void sm90_tf32_hc_prenorm_gemm(const torch::Tensor& a, smem_size = SM90ArchSpec::smem_capacity; + // Compile + const auto kernel = jit->compile("sm90_tf32_hc_prenorm_gemm", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm90_tf32_hc_prenorm_gemm_impl< + {}, {}, + {}, {}, {}, + {}, + {}, + {}, + {}, {} + >); +}}; +)", + n, k, + block_m, block_n, block_k, + num_splits, + swizzle_cd_mode, + num_stages, + num_math_threads, num_tma_threads)); + // Launch - const SM90BF16HCPrenormGemmRuntime::Args& args = { - .m = m, .n = n, .k = k, - .block_m = block_m, .block_n = block_n, .block_k = block_k, - .num_splits = num_splits, - .swizzle_cd_mode = swizzle_cd_mode, - .num_stages = num_stages, - .num_math_threads = num_math_threads, - .num_tma_threads = num_tma_threads, - .launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_threads, smem_size), - .tensor_map_a = tensor_map_a, - .tensor_map_b = tensor_map_b, - .tensor_map_d = tensor_map_d, - .sqr_sum = sqr_sum.data_ptr() - }; - const auto code = SM90BF16HCPrenormGemmRuntime::generate(args); - const auto runtime = compiler->build("sm90_tf32_hc_prenorm_gemm", code); - SM90BF16HCPrenormGemmRuntime::launch(runtime, args); + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_splits * ceil_div(m, block_m), 1, 1), + .block_dim = dim3(num_threads, 1, 1), + }, + m, tensor_map_a, tensor_map_b, tensor_map_d, sqr_sum.data_ptr() + ); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/smxx_clean_logits.hpp b/csrc/jit_kernels/impls/smxx_clean_logits.hpp deleted file mode 100644 index ebe4c7a6b9..0000000000 --- a/csrc/jit_kernels/impls/smxx_clean_logits.hpp +++ /dev/null @@ -1,81 +0,0 @@ -#pragma once - -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" -#include "../../utils/exception.hpp" - -namespace deep_gemm { - -class SMXXCleanLogitsRuntime final: public LaunchRuntime { -public: - struct Args { - int next_n; - int seq_len; - int seq_len_kv; - uint64_t stride_logits; - - int* cu_seq_len_k_start; - int* cu_seq_len_k_end; - void* logits; - at::ScalarType logits_dtype; - - int block_kv; - int num_warps; - - 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(&smxx_clean_logits< - {}, {}, {}, {} - >); -}}; -)", args.next_n, args.block_kv, args.num_warps, to_string(args.logits_dtype)); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.seq_len, args.seq_len_kv, static_cast(args.stride_logits), - args.cu_seq_len_k_start, args.cu_seq_len_k_end, args.logits - )); - } -}; - -static void smxx_clean_logits(const torch::Tensor& logits, - const std::optional& cu_seq_len_k_start, - const torch::Tensor& cu_seq_len_k_end, - const int& next_n, - const int& seq_len, const int& seq_len_kv, - const uint64_t &stride_logits) { - const int block_kv = 8192; - const int num_warps = 8; - const int smem_size = block_kv * sizeof(float); - - // Launch - const SMXXCleanLogitsRuntime::Args& args = { - .next_n = next_n, - .seq_len = seq_len, - .seq_len_kv = seq_len_kv, - .stride_logits = stride_logits, - .cu_seq_len_k_start = cu_seq_len_k_start.has_value() ? cu_seq_len_k_start.value().data_ptr() : nullptr, - .cu_seq_len_k_end = cu_seq_len_k_end.data_ptr(), - .logits = logits.data_ptr(), - .logits_dtype = logits.scalar_type(), - .block_kv = block_kv, - .num_warps = num_warps, - .launch_args = LaunchArgs(device_runtime->get_num_sms(), - num_warps * 32, smem_size) - }; - const auto code = SMXXCleanLogitsRuntime::generate(args); - const auto runtime = compiler->build("smxx_clean_logits", code); - SMXXCleanLogitsRuntime::launch(runtime, args); -} - -} // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/smxx_cublaslt.hpp b/csrc/jit_kernels/impls/smxx_cublaslt.hpp index 7f29b0a59a..181ed3edc0 100644 --- a/csrc/jit_kernels/impls/smxx_cublaslt.hpp +++ b/csrc/jit_kernels/impls/smxx_cublaslt.hpp @@ -5,23 +5,22 @@ #include #include -#include "../../jit/device_runtime.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" #include "../../utils/compatibility.hpp" namespace deep_gemm { -static auto get_cublaslt_layout(const cudaDataType& type, const int& rows, const int& cols, const int& ld, +static auto get_cublaslt_layout(const cudaDataType& type, const int& rows, const int& cols, const int64_t& ld, const std::optional& batch_count = std::nullopt, - const std::optional& batch_offset = std::nullopt) { + const std::optional& batch_offset = std::nullopt) { cublasLtMatrixLayout_t layout; DG_CUBLASLT_CHECK(cublasLtMatrixLayoutCreate(&layout, type, rows, cols, ld)); if (batch_count.has_value()) { DG_HOST_ASSERT(batch_offset.has_value()); - const int64_t batch_offset_int64 = batch_offset.value(); DG_CUBLASLT_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT, &batch_count.value(), sizeof(batch_count.value()))); - DG_CUBLASLT_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &batch_offset_int64, sizeof(batch_offset_int64))); + DG_CUBLASLT_CHECK(cublasLtMatrixLayoutSetAttribute(layout, CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET, &batch_offset.value(), sizeof(batch_offset.value()))); } return layout; } @@ -34,8 +33,17 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a, const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, - const bool& accumulate) { - cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F_FAST_TF32; + const bool& accumulate, + const float& alpha = 1.0f, + // NOTES: block-scaled UE4M3 scale pointers for the cuBLASLt A/B operands + // (i.e. the already-swapped column-major operands), `nullptr` means no block scaling + const void* a_block_sf = nullptr, + const void* b_block_sf = nullptr) { + const bool with_block_sf = a_block_sf != nullptr; + DG_HOST_ASSERT(with_block_sf == (b_block_sf != nullptr)); + + // Block-scaled matmuls require the plain FP32 compute type + cublasComputeType_t compute_type = with_block_sf ? CUBLAS_COMPUTE_32F : CUBLAS_COMPUTE_32F_FAST_TF32; cudaDataType_t scale_type = CUDA_R_32F; // Operation description @@ -45,22 +53,28 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a, DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b))); DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type))); -#if DG_CUBLASLT_ADVANCED_FEATURES_COMPATIBLE - const int math_sms = device_runtime->get_num_sms(); - DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &math_sms, sizeof(math_sms))); -#endif + if (with_block_sf) { + // NOTES: cuBLASLt only supports the NVFP4 recipe for FP4 operands + // (UE4M3 scaling factors with 16-element blocks) + const int32_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, &scale_mode, sizeof(scale_mode))); + DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, &scale_mode, sizeof(scale_mode))); + DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, &a_block_sf, sizeof(a_block_sf))); + DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, &b_block_sf, sizeof(b_block_sf))); + } + + const int num_math_sms = runtime->get_num_sms(); + DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &num_math_sms, sizeof(num_math_sms))); -#if DG_FP8_COMPATIBLE and DG_CUBLASLT_ADVANCED_FEATURES_COMPATIBLE bool fp8_fast_accumulate = false; if (a.scalar_type() == torch::kFloat8_e4m3fn) DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_FAST_ACCUM, &fp8_fast_accumulate, sizeof(fp8_fast_accumulate))); -#endif // Get cuBLASLt handle, workspace, and stream - const auto handle = device_runtime->get_cublaslt_handle(); - const auto workspace = device_runtime->get_cublaslt_workspace(); - const auto workspace_bytes = workspace.nbytes(); + const auto handle = runtime->get_cublaslt_handle(); const auto stream = at::cuda::getCurrentCUDAStream(); + const auto workspace = runtime->get_cublaslt_workspace(stream); + const auto workspace_bytes = workspace.nbytes(); // Algorithm selection cublasLtMatmulPreference_t pref; @@ -77,7 +91,7 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a, DG_HOST_ASSERT(num_heuristic_results == 1 and "Unable to find any algorithm for the GEMM"); // Call: D = alpha * (A @ B) + beta * C - const float alpha = 1.0, beta = accumulate ? 1.0 : 0.0; + const float beta = accumulate ? 1.0f : 0.0f; DG_CUBLASLT_CHECK(cublasLtMatmul(handle, // Light handle desc, // Operation description &alpha, // Alpha @@ -102,7 +116,8 @@ static void cublaslt_gemm(const torch::Tensor& lhs, const torch::Tensor& rhs, const torch::Tensor& out, const int& m, const int& n, const int& k, const cute::UMMA::Major& a_major, const cute::UMMA::Major& b_major, - const bool& accumulate) { + const bool& accumulate, + const float& alpha = 1.0f) { const auto trans_a = b_major == cute::UMMA::Major::K ? CUBLAS_OP_T : CUBLAS_OP_N; const auto trans_b = a_major == cute::UMMA::Major::K ? CUBLAS_OP_N : CUBLAS_OP_T; @@ -116,7 +131,57 @@ static void cublaslt_gemm(const torch::Tensor& lhs, const torch::Tensor& rhs, : get_cublaslt_layout(cuda_type_b, m, k, lhs.stride(1)); const auto layout_d = get_cublaslt_layout(cuda_type_d, n, m, out.stride(0)); - call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, accumulate); + call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, accumulate, alpha); +} + +static void cublaslt_batched_gemm(const torch::Tensor& lhs, const torch::Tensor& rhs, + const torch::Tensor& out, + const uint32_t& m, const uint32_t& n, const uint32_t& k, + const uint32_t& num_batches, + const cute::UMMA::Major& a_major, const cute::UMMA::Major& b_major) { + const auto trans_a = b_major == cute::UMMA::Major::K ? CUBLAS_OP_T : CUBLAS_OP_N; + const auto trans_b = a_major == cute::UMMA::Major::K ? CUBLAS_OP_N : CUBLAS_OP_T; + const auto lhs_batch_offset = lhs.stride(0); + const auto rhs_batch_offset = rhs.stride(0); + const auto out_batch_offset = out.stride(0); + + // cuBLASLt uses column-major layouts. Swap the operands so that their column-major + // interpretation computes the transpose of the requested row-major result. + const auto cuda_type_a = at::cuda::ScalarTypeToCudaDataType(rhs.scalar_type()); + const auto cuda_type_b = at::cuda::ScalarTypeToCudaDataType(lhs.scalar_type()); + const auto cuda_type_d = at::cuda::ScalarTypeToCudaDataType(out.scalar_type()); + const auto layout_a = b_major == cute::UMMA::Major::K ? + get_cublaslt_layout(cuda_type_a, k, n, rhs.stride(-2), num_batches, rhs_batch_offset) : + get_cublaslt_layout(cuda_type_a, n, k, rhs.stride(-1), num_batches, rhs_batch_offset); + const auto layout_b = a_major == cute::UMMA::Major::K ? + get_cublaslt_layout(cuda_type_b, k, m, lhs.stride(-2), num_batches, lhs_batch_offset) : + get_cublaslt_layout(cuda_type_b, m, k, lhs.stride(-1), num_batches, lhs_batch_offset); + const auto layout_d = get_cublaslt_layout(cuda_type_d, n, m, out.stride(-2), + num_batches, out_batch_offset); + + call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, false); +} + +static void cublaslt_nvfp4_gemm(const torch::Tensor& lhs, const torch::Tensor& lhs_sf, + const torch::Tensor& rhs, const torch::Tensor& rhs_sf, + const torch::Tensor& out, + const int& m, const int& n, const int& k, + const bool& accumulate) { + // NOTES: block-scaled FP4 only supports the NT layout (both operands K-major), + // which maps to the column-major TN layout required by cuBLASLt + const auto trans_a = CUBLAS_OP_T; + const auto trans_b = CUBLAS_OP_N; + + // Matrix layouts + // NOTES: dimensions and leading dims are in logical FP4 elements (2 elements per packed byte) + const auto cuda_type_d = at::cuda::ScalarTypeToCudaDataType(out.scalar_type()); + const auto layout_a = get_cublaslt_layout(CUDA_R_4F_E2M1, k, n, static_cast(rhs.stride(0)) * 2); + const auto layout_b = get_cublaslt_layout(CUDA_R_4F_E2M1, k, m, static_cast(lhs.stride(0)) * 2); + const auto layout_d = get_cublaslt_layout(cuda_type_d, n, m, static_cast(out.stride(0))); + + // NOTES: after the column-major operand swap, the cuBLASLt A operand is `rhs` + call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, accumulate, + 1.0f, rhs_sf.data_ptr(), lhs_sf.data_ptr()); } static void cublaslt_bhr_hdr_bhd(const torch::Tensor& lhs, const torch::Tensor& rhs, const torch::Tensor& out, @@ -148,4 +213,19 @@ static void cublaslt_bhd_hdr_bhr(const torch::Tensor& lhs, const torch::Tensor& call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, false); } +static void cublaslt_bhd_bhr_hdr(const torch::Tensor& lhs, const torch::Tensor& rhs, const torch::Tensor& out, + const uint32_t& b, const uint32_t& h, const uint32_t& r, const uint32_t& d, + const bool& accumulate) { + const auto m = r, n = d, k = b; + const auto trans_a = CUBLAS_OP_N; + const auto trans_b = CUBLAS_OP_T; + + // Matrix layouts + const auto layout_a = get_cublaslt_layout(CUDA_R_16BF, m, k, rhs.stride(0), h, rhs.stride(1)); + const auto layout_b = get_cublaslt_layout(CUDA_R_16BF, n, k, lhs.stride(0), h, lhs.stride(1)); + const auto layout_d = get_cublaslt_layout(CUDA_R_32F, m, n, out.stride(1), h, out.stride(0)); + + call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, accumulate); +} + } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/smxx_layout.hpp b/csrc/jit_kernels/impls/smxx_layout.hpp index 46374f274e..a2b651bda8 100644 --- a/csrc/jit_kernels/impls/smxx_layout.hpp +++ b/csrc/jit_kernels/impls/smxx_layout.hpp @@ -1,80 +1,17 @@ #pragma once +#include #include -#include "../../jit/kernel_runtime.hpp" -#include "../../jit/compiler.hpp" +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" -#include "../../utils/format.hpp" #include "../../utils/math.hpp" #include "../../utils/layout.hpp" #include "../heuristics/runtime.hpp" namespace deep_gemm { -class TransposeFP32Runtime final: public LaunchRuntime { -public: - struct Args { - int mn, sf_k; - int block_mn; - void *sf, *out; - - 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(&transpose_fp32< - {}, {}, {} - >); -}}; -)", args.launch_args.num_threads, args.block_mn, args.sf_k); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, args.sf, args.out, static_cast(args.mn))); - } -}; - -class TransposeAndPackFP32IntoUE8M0Runtime final: public LaunchRuntime { -public: - struct Args { - int mn, sf_k; - int num_psum_groups, m_alignment; - bool use_psum_layout; - int block_mn; - void *sf, *out, *grouped_layout; - - 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(&transpose_and_pack_fp32_into_ue8m0< - {}, {}, {}, {}, {} - >); -}}; -)", args.launch_args.num_threads, args.block_mn, args.sf_k, - args.num_psum_groups, args.use_psum_layout ? "true" : "false"); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, args.sf, args.out, static_cast(args.mn), - args.grouped_layout, static_cast(args.m_alignment))); - } -}; - -class PackFP32IntoUE8M0Runtime final: public LaunchRuntime { +class PackFP32IntoUE8M0Runtime final { public: struct Args { int num_groups, mn, sf_k, packed_sf_k, gran_k, k_alignment; @@ -82,11 +19,11 @@ class PackFP32IntoUE8M0Runtime final: public LaunchRuntimecompile(tag, std::format(R"( #include using namespace deep_gemm; @@ -96,12 +33,15 @@ static void __instantiate_kernel() {{ {}, {}, {}, {}, {}, {} >); }}; -)", args.num_groups, args.launch_args.num_threads, args.block_mn, args.block_packed_sf_k, "true", args.use_psum_layout ? "true" : "false"); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.sf, args.out, args.grouped_layout, args.mn, args.sf_k, args.packed_sf_k, args.gran_k, args.k_alignment)); +)", args.num_groups, args.options.block_dim->x, args.block_mn, args.block_packed_sf_k, + "true", args.use_psum_layout ? "true" : "false")); + + // Launch + jit->launch( + kernel, args.options, + args.sf, args.out, args.grouped_layout, args.mn, args.sf_k, args.packed_sf_k, + args.gran_k, args.k_alignment + ); } }; @@ -134,20 +74,33 @@ static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) { out.copy_(batched_sf); } else { constexpr int block_mn = 64; + constexpr int block_sf_k = 128; constexpr int num_threads = 512; - const auto smem_size = block_mn * (sf_k + (1 - (sf_k % 2))) * static_cast(sizeof(float)); - const TransposeFP32Runtime::Args& args = { - .mn = mn, - .sf_k = sf_k, - .block_mn = block_mn, - .sf = batched_sf.data_ptr(), - .out = out.data_ptr(), - .launch_args = LaunchArgs({ceil_div(mn, block_mn), num_sf_batches}, num_threads, smem_size) - }; + const auto smem_sf_k = sf_k < block_sf_k ? sf_k : block_sf_k; + const auto smem_size = block_mn * (smem_sf_k + (1 - smem_sf_k % 2)) * static_cast(sizeof(float)); + + // Compile + const auto kernel = jit->compile("transpose_fp32", std::format(R"( +#include - const auto code = TransposeFP32Runtime::generate(args); - const auto runtime = compiler->build("transpose_fp32", code); - TransposeFP32Runtime::launch(runtime, args); +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&transpose_fp32< + {}, {}, {}, {} + >); +}}; +)", num_threads, block_mn, sf_k, block_sf_k)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(ceil_div(mn, block_mn) * ceil_div(sf_k, block_sf_k), num_sf_batches, 1), + .block_dim = dim3(num_threads, 1, 1), + }, + batched_sf.data_ptr(), out.data_ptr(), static_cast(mn) + ); } return (dim == 2) ? out.squeeze(0) : out; } @@ -195,31 +148,40 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T const auto m_alignment = use_psum_layout ? heuristics_runtime->get_mk_alignment_for_contiguous_layout() : 0; const auto num_psum_groups = use_psum_layout ? static_cast(psum_layout->numel()) : 1; - // Launch the kernel if (batched_sf.is_contiguous()) { if ((mn * sf_k) % 4 != 0 and num_sf_batches > 1) return get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(sf); constexpr int block_mn = 48; + constexpr int block_sf_k = 128; constexpr int num_threads = 512; const auto psum_smem_elems = use_psum_layout ? align(num_psum_groups * 2, 4) : 0; - const auto smem_size = block_mn * sf_k * 4 + psum_smem_elems * 4; - const TransposeAndPackFP32IntoUE8M0Runtime::Args& args = { - .mn = mn, - .sf_k = sf_k, - .num_psum_groups = num_psum_groups, - .m_alignment = m_alignment, - .use_psum_layout = use_psum_layout, - .block_mn = block_mn, - .sf = batched_sf.data_ptr(), - .out = out.data_ptr(), - .grouped_layout = use_psum_layout ? psum_layout->data_ptr() : nullptr, - .launch_args = LaunchArgs({ceil_div(mn, block_mn), num_sf_batches}, num_threads, smem_size) - }; + const auto smem_sf_k = sf_k < block_sf_k ? sf_k : block_sf_k + 1; + const auto smem_size = block_mn * smem_sf_k * 4 + psum_smem_elems * 4; - const auto code = TransposeAndPackFP32IntoUE8M0Runtime::generate(args); - const auto runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code); - TransposeAndPackFP32IntoUE8M0Runtime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("transpose_and_pack_fp32_into_ue8m0", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&transpose_and_pack_fp32_into_ue8m0< + {}, {}, {}, {}, {}, {} + >); +}}; +)", num_threads, block_mn, sf_k, block_sf_k, num_psum_groups, use_psum_layout ? "true" : "false")); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(ceil_div(mn, block_mn) * ceil_div(sf_k, block_sf_k), num_sf_batches, 1), + .block_dim = dim3(num_threads, 1, 1), + }, + batched_sf.data_ptr(), out.data_ptr(), static_cast(mn), + use_psum_layout ? psum_layout->data_ptr() : nullptr, static_cast(m_alignment) + ); } else { DG_HOST_ASSERT(not use_psum_layout); if (mn % 4 != 0 or num_sf_batches > 1) @@ -229,25 +191,26 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T constexpr int block_mn = 128; constexpr int block_packed_sf_k = 16; constexpr int num_threads = 512; - const PackFP32IntoUE8M0Runtime::Args& args = { + + // Compile and launch + PackFP32IntoUE8M0Runtime::compile_and_launch("pack_fp32_into_ue8m0", { .num_groups = 1, .mn = mn, .sf_k = sf_k, .packed_sf_k = packed_sf_k, - // NOTES: `gran_k`/`k_alignment` are only used by k-grouped psum, dummy here .gran_k = 128, .k_alignment = 128, + .use_psum_layout = false, .block_mn = block_mn, .block_packed_sf_k = block_packed_sf_k, .sf = batched_sf.data_ptr(), .out = out.data_ptr(), .grouped_layout = nullptr, - .launch_args = LaunchArgs({ceil_div(mn, block_mn), ceil_div(packed_sf_k, block_packed_sf_k)}, num_threads) - }; - - const auto code = PackFP32IntoUE8M0Runtime::generate(args); - const auto runtime = compiler->build("pack_fp32_into_ue8m0", code); - PackFP32IntoUE8M0Runtime::launch(runtime, args); + .options = { + .grid_dim = dim3(ceil_div(mn, block_mn), ceil_div(packed_sf_k, block_packed_sf_k), 1), + .block_dim = dim3(num_threads, 1, 1), + } + }); } return (dim == 2) ? out.squeeze(0) : out; } @@ -258,8 +221,7 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons const int gran_k, const int k_alignment, const bool& use_psum_layout) { - DG_HOST_ASSERT(gran_k == 32 or gran_k == 128); - DG_HOST_ASSERT(k_alignment % 32 == 0); + DG_HOST_ASSERT(jit->device.get_arch_major() == 10 and (gran_k == 32 or gran_k == 128) and k_alignment % 128 == 0); const auto [sf_k, mn] = get_shape<2>(sf); const auto num_groups = static_cast(grouped_layout.numel()); @@ -278,11 +240,15 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons if (has_synced_ks) { int ref_sf_k = 0; for (const auto k: ks_cpu.value()) { - ref_sf_k += ceil_div(k, gran_k); - packed_sf_k += ceil_div(k, gran_k * 4); + DG_HOST_ASSERT(k % k_alignment == 0 and k % gran_k == 0); + const auto group_sf_k = k / gran_k; + ref_sf_k += group_sf_k; + packed_sf_k += ceil_div(group_sf_k, 4); } - DG_HOST_ASSERT(use_psum_layout or ref_sf_k == sf_k); + DG_HOST_ASSERT(ref_sf_k == sf_k); } else { + // Exact group sizes are read from the psum layout by the pack kernel. + // This upper bound allows three tail-padding slots per group. packed_sf_k = (sf_k + num_groups * 3) / 4; } if (packed_sf_k == 0) @@ -293,7 +259,9 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons constexpr int block_mn = 128; constexpr int block_packed_sf_k = 16; constexpr int num_threads = 512; - const PackFP32IntoUE8M0Runtime::Args args = { + + // Compile and launch + PackFP32IntoUE8M0Runtime::compile_and_launch("pack_fp32_into_ue8m0", { .num_groups = num_groups, .mn = mn, .sf_k = sf_k, @@ -306,12 +274,11 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons .sf = sf.data_ptr(), .out = out.data_ptr(), .grouped_layout = grouped_layout.data_ptr(), - .launch_args = LaunchArgs({ceil_div(mn, block_mn), ceil_div(packed_sf_k, block_packed_sf_k)}, num_threads) - }; - - const auto code = PackFP32IntoUE8M0Runtime::generate(args); - const auto runtime = compiler->build("pack_fp32_into_ue8m0", code); - PackFP32IntoUE8M0Runtime::launch(runtime, args); + .options = { + .grid_dim = dim3(ceil_div(mn, block_mn), ceil_div(packed_sf_k, block_packed_sf_k), 1), + .block_dim = dim3(num_threads, 1, 1), + } + }); return out; } @@ -322,9 +289,8 @@ static torch::Tensor check_k_grouped_packed_ue8m0_tensor(const torch::Tensor& sf const int gran_k, const int k_alignment, const bool& use_psum_layout) { - DG_HOST_ASSERT(gran_k == 32 or gran_k == 128); - DG_HOST_ASSERT(k_alignment % 32 == 0); DG_HOST_ASSERT(sf.scalar_type() == torch::kInt); + DG_HOST_ASSERT(gran_k == 32); DG_HOST_ASSERT(sf.dim() == 2); DG_HOST_ASSERT(sf.is_contiguous()); @@ -338,12 +304,12 @@ static torch::Tensor check_k_grouped_packed_ue8m0_tensor(const torch::Tensor& sf const auto has_synced_ks = ks_cpu.has_value() and not ks_cpu.value().empty(); if (has_synced_ks) { DG_HOST_ASSERT(static_cast(ks_cpu.value().size()) == num_groups); - if (not use_psum_layout) { - int aligned_packed_sf_k = 0; - for (const auto k: ks_cpu.value()) - aligned_packed_sf_k += ceil_div(k, gran_k * 4); - DG_HOST_ASSERT(packed_sf_k >= aligned_packed_sf_k); + int required_packed_sf_k = 0; + for (const auto k: ks_cpu.value()) { + DG_HOST_ASSERT(k % k_alignment == 0); + required_packed_sf_k += ceil_div(k, gran_k * 4); } + DG_HOST_ASSERT(packed_sf_k >= required_packed_sf_k); } else { DG_HOST_ASSERT(use_psum_layout); } diff --git a/csrc/python_api.cpp b/csrc/python_api.cpp index a966afe1ed..b1de15c4a6 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -1,13 +1,18 @@ #include #include +#include +#include + +#include "apis/config.hpp" #include "apis/attention.hpp" #include "apis/einsum.hpp" #include "apis/hyperconnection.hpp" #include "apis/gemm.hpp" #include "apis/layout.hpp" -#include "apis/mega.hpp" -#include "apis/runtime.hpp" +#include "apis/mega_moe.hpp" +#include "apis/mega_mhc.hpp" +#include "apis/mega_gate.hpp" #ifndef TORCH_EXTENSION_NAME #define TORCH_EXTENSION_NAME _C @@ -17,6 +22,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.doc() = "DeepGEMM C++ library"; + // Register JIT objects + deep_jit::register_python_api(m, deep_gemm::jit); + + // Register config APIs + deep_gemm::config::register_apis(m); + + // Register kernels // TODO: make SM80 incompatible issues raise errors deep_gemm::attention::register_apis(m); deep_gemm::einsum::register_apis(m); @@ -24,5 +36,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { deep_gemm::gemm::register_apis(m); deep_gemm::layout::register_apis(m); deep_gemm::mega::register_apis(m); - deep_gemm::runtime::register_apis(m); + deep_gemm::mega_mhc::register_apis(m); + deep_gemm::mega_gate::register_apis(m); } diff --git a/csrc/runtime/jit.hpp b/csrc/runtime/jit.hpp new file mode 100644 index 0000000000..0a394de2bb --- /dev/null +++ b/csrc/runtime/jit.hpp @@ -0,0 +1,34 @@ +#pragma once + +#include +#include +#include + +#include +#include + +namespace deep_gemm { + +inline deep_jit::LazyInit> jit(nullptr); + +inline void init_jit(const std::string& library_root_path) { + const auto library_root = std::filesystem::absolute(library_root_path); + const auto include_dir = library_root / "include"; + const auto config = deep_jit::Config( + library_root, + "DG", + "cutlass-" + std::to_string(CUTLASS_VERSION), + {include_dir}, + {"deep_gemm/"}); + + jit = deep_jit::LazyInit>([config] { + auto runtime = std::make_shared>(config); + runtime->default_compiler_options.nvcc_flags-> + emplace_back("--diag-suppress=39,161,174,177,186,940"); + runtime->default_compiler_options.nvcc_flags-> + emplace_back("--compiler-options=-Wno-deprecated-declarations,-Wno-abi"); + return runtime; + }); +} + +} // namespace deep_gemm diff --git a/csrc/runtime/runtime.hpp b/csrc/runtime/runtime.hpp new file mode 100644 index 0000000000..1eb53f1aaf --- /dev/null +++ b/csrc/runtime/runtime.hpp @@ -0,0 +1,99 @@ +#pragma once + +#include +#include +#include + +#include +#include + +#include +#include + +#include "../utils/exception.hpp" +#include "jit.hpp" + +namespace deep_gemm { + +class Runtime { +public: + int num_sms = 0, tc_util = 0; + + // cuBLASLt utils + // cuBLAS will select worse heuristics when workspace > 16 MiB with shape, e.g. m=128, n=7168, k=16384 + static constexpr size_t kCublasLtWorkspaceSize = 16 * 1024 * 1024; + // cuBLASLt may use the workspace asynchronously, so concurrent streams cannot share it. + std::unordered_map cublaslt_workspaces; + + // Create the cuBLASLt handle ourselves + cublasLtHandle_t cublaslt_handle; + bool use_pytorch_managed_cublaslt_handle; + bool use_temp_cublaslt_workspace; + + explicit Runtime() { + // Whether to use PyTorch cuBLASLt + // By default, we don't use it, + // as `at::cuda::getCurrentCUDABlasLtHandle` has large CPU overhead with some PyTorch versions + use_pytorch_managed_cublaslt_handle = deep_jit::get_env("DG_USE_PYTORCH_CUBLASLT_HANDLE", 0) > 0; + // Whether to create workspace tensor on each call instead of holding one. + // Enabled by compute-sanitizer tests, which trigger `cudaErrorCudartUnloading` + // when the workspace tensor is destructed after CUDA driver shutdown. + use_temp_cublaslt_workspace = deep_jit::get_env("DG_USE_TEMP_CUBLASLT_WORKSPACE", 0) > 0; + + if (not use_pytorch_managed_cublaslt_handle) + DG_CUBLASLT_CHECK(cublasLtCreate(&cublaslt_handle)); + } + + ~Runtime() noexcept(false) { + if (not use_pytorch_managed_cublaslt_handle) + DG_CUBLASLT_CHECK(cublasLtDestroy(cublaslt_handle)); + } + + cublasLtHandle_t get_cublaslt_handle() const { + if (use_pytorch_managed_cublaslt_handle) + return at::cuda::getCurrentCUDABlasLtHandle(); + + // Self-managed handle + return cublaslt_handle; + } + + torch::Tensor get_cublaslt_workspace(const c10::cuda::CUDAStream& stream) { + const auto options = dtype(torch::kByte).device(stream.device()); + if (use_temp_cublaslt_workspace) + return torch::empty({kCublasLtWorkspaceSize}, options); + + auto& workspace = cublaslt_workspaces[stream]; + if (not workspace.defined()) + workspace = torch::empty({kCublasLtWorkspaceSize}, options); + return workspace; + } + + void set_num_sms(const int& new_num_sms) { + DG_HOST_ASSERT(0 < new_num_sms and new_num_sms <= jit->device.get_num_sms()); + num_sms = new_num_sms; + } + + int get_num_sms() { + if (num_sms == 0) + num_sms = jit->device.get_num_sms(); + return num_sms; + } + + bool is_cublaslt_available() { + return get_num_sms() == jit->device.get_num_sms(); + } + + void set_tc_util(const int& new_tc_util) { + DG_HOST_ASSERT(0 <= new_tc_util and new_tc_util <= 100); + tc_util = new_tc_util; + } + + int get_tc_util() const { + return tc_util == 0 ? 100 : tc_util; + } + +}; + +inline auto runtime = deep_jit::LazyInit([](){ return std::make_shared(); }); + +} // namespace deep_gemm diff --git a/csrc/utils/compatibility.hpp b/csrc/utils/compatibility.hpp index 9e2d67205a..a327753200 100644 --- a/csrc/utils/compatibility.hpp +++ b/csrc/utils/compatibility.hpp @@ -4,14 +4,9 @@ #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)) +#include -// `cuTensorMapEncodeTiled` is supported since CUDA Driver API 12.1 -#define DG_TENSORMAP_COMPATIBLE (CUDA_VERSION >= 12010) - -// `cublasGetErrorString` is supported since CUDA Runtime API 11.4.2 -#define DG_CUBLAS_GET_ERROR_STRING_COMPATIBLE (CUDART_VERSION >= 11042) - -// `CUBLASLT_MATMUL_DESC_FAST_ACCUM` and `CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET` are supported since CUDA Runtime API 11.8 -#define DG_CUBLASLT_ADVANCED_FEATURES_COMPATIBLE (CUDART_VERSION >= 11080) \ No newline at end of file +DG_STATIC_ASSERT(TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 3), + "DeepGEMM requires PyTorch 2.3 or newer"); +DG_STATIC_ASSERT(CUDA_VERSION >= 12090, "DeepGEMM requires CUDA Driver API 12.9 or newer"); +DG_STATIC_ASSERT(CUDART_VERSION >= 12090, "DeepGEMM requires CUDA Runtime API 12.9 or newer"); diff --git a/csrc/utils/exception.hpp b/csrc/utils/exception.hpp index 417dd3b42f..8b5eb0b795 100644 --- a/csrc/utils/exception.hpp +++ b/csrc/utils/exception.hpp @@ -1,9 +1,10 @@ #pragma once -#include #include -#include #include +#include + +#include #include "compatibility.hpp" @@ -39,30 +40,6 @@ do { \ #define DG_HOST_UNREACHABLE(reason) (throw DGException("Assertion", __FILE__, __LINE__, reason)) #endif -#ifndef DG_NVRTC_CHECK -#define DG_NVRTC_CHECK(cmd) \ -do { \ - const auto e = (cmd); \ - if (e != NVRTC_SUCCESS) { \ - throw DGException("NVRTC", __FILE__, __LINE__, nvrtcGetErrorString(e)); \ - } \ -} while (0) -#endif - -#ifndef DG_CUDA_DRIVER_CHECK -#define DG_CUDA_DRIVER_CHECK(cmd) \ -do { \ - const auto e = (cmd); \ - if (e != CUDA_SUCCESS) { \ - std::stringstream ss; \ - const char *name, *info; \ - lazy_cuGetErrorName(e, &name), lazy_cuGetErrorString(e, &info); \ - ss << static_cast(e) << " (" << name << ", " << info << ")"; \ - throw DGException("CUDA driver", __FILE__, __LINE__, ss.str()); \ - } \ -} while (0) -#endif - #ifndef DG_CUDA_RUNTIME_CHECK #define DG_CUDA_RUNTIME_CHECK(cmd) \ do { \ @@ -76,25 +53,6 @@ do { \ #endif #ifndef DG_CUBLASLT_CHECK - -#if !DG_CUBLAS_GET_ERROR_STRING_COMPATIBLE -inline const char* cublasGetStatusString(cublasStatus_t status) { - switch(status) { - case CUBLAS_STATUS_SUCCESS: return "CUBLAS_STATUS_SUCCESS"; - case CUBLAS_STATUS_NOT_INITIALIZED: return "CUBLAS_STATUS_NOT_INITIALIZED"; - case CUBLAS_STATUS_ALLOC_FAILED: return "CUBLAS_STATUS_ALLOC_FAILED"; - case CUBLAS_STATUS_INVALID_VALUE: return "CUBLAS_STATUS_INVALID_VALUE"; - case CUBLAS_STATUS_ARCH_MISMATCH: return "CUBLAS_STATUS_ARCH_MISMATCH"; - case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR"; - case CUBLAS_STATUS_EXECUTION_FAILED: return "CUBLAS_STATUS_EXECUTION_FAILED"; - case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR"; - case CUBLAS_STATUS_NOT_SUPPORTED: return "CUBLAS_STATUS_NOT_SUPPORTED"; - case CUBLAS_STATUS_LICENSE_ERROR: return "CUBLAS_STATUS_LICENSE_ERROR"; - default: return "Unknown cuBLAS error"; - } -} -#endif - #define DG_CUBLASLT_CHECK(cmd) \ do { \ const auto e = (cmd); \ diff --git a/csrc/utils/format.hpp b/csrc/utils/format.hpp deleted file mode 100644 index bf617372be..0000000000 --- a/csrc/utils/format.hpp +++ /dev/null @@ -1,6 +0,0 @@ -#pragma once - -// Just a wrapper for the `fmt` headers -#define FMT_HEADER_ONLY -#include -#include diff --git a/csrc/utils/hash.hpp b/csrc/utils/hash.hpp deleted file mode 100644 index 9efe640816..0000000000 --- a/csrc/utils/hash.hpp +++ /dev/null @@ -1,39 +0,0 @@ -#pragma once - -#include - -namespace deep_gemm { - -static uint64_t fnv1a(const std::vector& data, const uint64_t& seed) { - uint64_t h = seed; - const uint64_t prime = 0x100000001b3ull; - for (const char& c: data) { - h ^= static_cast(c); - h *= prime; - } - return h; -} - -static std::string get_hex_digest(const std::vector& data) { - const auto state_0 = fnv1a(data, 0xc6a4a7935bd1e995ull); - const auto state_1 = fnv1a(data, 0x9e3779b97f4a7c15ull); - - // Split-mix 64 - const auto split_mix = [](uint64_t z) { - z = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9ull; - z = (z ^ (z >> 27)) * 0x94d049bb133111ebull; - return z ^ (z >> 31); - }; - - std::ostringstream oss; - oss << std::hex << std::setfill('0') - << std::setw(16) << split_mix(state_0) - << std::setw(16) << split_mix(state_1); - return oss.str(); -} - -static std::string get_hex_digest(const std::string& data) { - return get_hex_digest(std::vector{data.begin(), data.end()}); -} - -} // namespace deep_gemm diff --git a/csrc/utils/layout.hpp b/csrc/utils/layout.hpp index 4707fc48e9..7a97cc6da8 100644 --- a/csrc/utils/layout.hpp +++ b/csrc/utils/layout.hpp @@ -5,32 +5,38 @@ #include "math.hpp" #include "exception.hpp" -#include "../jit/device_runtime.hpp" +#include "../runtime/jit.hpp" namespace deep_gemm { // Major-ness stuffs +template static void major_check(const torch::Tensor& t) { const auto dim = t.dim(); DG_HOST_ASSERT(dim == 2 or dim == 3); - if (dim == 3) - DG_HOST_ASSERT(t.stride(0) == t.size(-2) * t.size(-1)); + if constexpr (kRequireContiguousBatch) { + if (dim == 3) + DG_HOST_ASSERT(t.stride(0) == t.size(-2) * t.size(-1)); + } DG_HOST_ASSERT(t.stride(-2) == 1 or t.stride(-1) == 1); } +template static cute::UMMA::Major get_major_type_ab(const torch::Tensor& t) { - major_check(t); + major_check(t); return t.stride(-1) == 1 ? cute::UMMA::Major::K : cute::UMMA::Major::MN; } +template static void check_major_type_cd(const torch::Tensor& t) { // NOTES: the library only supports row-major output layouts - major_check(t); + major_check(t); DG_HOST_ASSERT(t.stride(-1) == 1); } -static bool fp8_requires_k_major() { - return device_runtime->get_arch_major() == 9; +static bool fp8_fp4_requires_k_major(const torch::Tensor& a, const torch::Tensor& b) { + return jit->device.get_arch_major() == 9 or + (a.scalar_type() == kPackedFP4 and b.scalar_type() == kPackedFP4); } // Tensor utils @@ -73,7 +79,7 @@ static std::tuple check_grouped_ab_fp8_fp4(const torch::Tensor& a // Recipe static std::tuple get_default_recipe(const torch::ScalarType& sfa_dtype, const torch::ScalarType& sfb_dtype) { - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major == 9) { DG_HOST_ASSERT(sfa_dtype == torch::kFloat and sfb_dtype == torch::kFloat); return {1, 128, 128}; diff --git a/csrc/utils/lazy_init.hpp b/csrc/utils/lazy_init.hpp deleted file mode 100644 index 386b1b4539..0000000000 --- a/csrc/utils/lazy_init.hpp +++ /dev/null @@ -1,27 +0,0 @@ -#pragma once - -#include -#include - -#define DG_DECLARE_STATIC_VAR_IN_CLASS(cls, name) decltype(cls::name) cls::name - -namespace deep_gemm { - -template -class LazyInit { -public: - explicit LazyInit(std::function()> factory) - : factory(std::move(factory)) {} - - T* operator -> () { - if (ptr == nullptr) - ptr = factory(); - return ptr.get(); - } - -private: - std::shared_ptr ptr; - std::function()> factory; -}; - -} // namespace deep_gemm diff --git a/csrc/utils/system.hpp b/csrc/utils/system.hpp deleted file mode 100644 index fda020be55..0000000000 --- a/csrc/utils/system.hpp +++ /dev/null @@ -1,128 +0,0 @@ -#pragma once - -#include -#include -#include -#include -#include -#include -#include - -#include "exception.hpp" -#include "format.hpp" - -namespace deep_gemm { - -// ReSharper disable once CppNotAllPathsReturnValue -template -static dtype_t get_env(const std::string& name, const dtype_t& default_value = dtype_t()) { - const auto c_str = std::getenv(name.c_str()); - if (c_str == nullptr) - return default_value; - - // Read the env and convert to the desired type - if constexpr (std::is_same_v) { - return std::string(c_str); - } else if constexpr (std::is_same_v) { - int value; - std::sscanf(c_str, "%d", &value); - return value; - } else { - DG_HOST_ASSERT(false and "Unexpected type"); - } -} - -static std::tuple call_external_command(std::string command) { - command = command + " 2>&1"; - const auto deleter = [](FILE* f) { if (f) pclose(f); }; - std::unique_ptr pipe(popen(command.c_str(), "r"), deleter); - DG_HOST_ASSERT(pipe != nullptr); - - std::array buffer; - std::string output; - while (fgets(buffer.data(), buffer.size(), pipe.get())) - output += buffer.data(); - const auto status = pclose(pipe.release()); - // NOTES: if the child was killed by a signal (e.g., SIGINT from Ctrl+C), - // WEXITSTATUS would incorrectly return 0. Treat signal death as failure. - const auto exit_code = WIFEXITED(status) ? WEXITSTATUS(status) : 128 + WTERMSIG(status); - return {exit_code, output}; -} - -static std::vector collect_files(const std::filesystem::path& root) { - std::vector files; - std::function impl; - impl = [&](const std::filesystem::path& dir) { - for (const auto& entry: std::filesystem::directory_iterator(dir)) { - if (entry.is_directory()) { - impl(entry.path()); - } else if (entry.is_regular_file() and entry.path().extension() == ".cuh") { - files.emplace_back(entry.path()); - } - } - }; - impl(root); - - // Be consistent - std::sort(files.begin(), files.end()); - return files; -} - -static std::filesystem::path make_dirs(const std::filesystem::path& path) { - // OK if existed - std::error_code capture; - const bool created = std::filesystem::create_directories(path, capture); - if (not (created or capture.value() == 0)) { - DG_HOST_UNREACHABLE(fmt::format("Failed to make directory: {}, created: {}, value: {}", - path.c_str(), created, capture.value())); - } - if (created and get_env("DG_JIT_DEBUG")) - printf("Create directory: %s\n", path.c_str()); - return path; -} - -static std::string get_uuid() { - static std::random_device rd; - static std::mt19937 gen([]() { - return rd() ^ std::chrono::steady_clock::now().time_since_epoch().count(); - }()); - static std::uniform_int_distribution dist; - - std::stringstream ss; - ss << getpid() << "-" - << std::hex << std::setfill('0') - << std::setw(8) << dist(gen) << "-" - << std::setw(8) << dist(gen) << "-" - << std::setw(8) << dist(gen); - return ss.str(); -} - -static void safe_remove_all(const std::filesystem::path& path) { - std::error_code ec; - if (not std::filesystem::exists(path, ec) or ec) - return; - - // A single file - if (not std::filesystem::is_directory(path, ec) or ec) { - std::filesystem::remove(path, ec); - return; - } - - // Remove directory - auto it = std::filesystem::directory_iterator(path, - std::filesystem::directory_options::skip_permission_denied, ec); - for (auto end = std::filesystem::directory_iterator(); it != end and not ec;) { - const auto entry_path = it->path(); - - // Increase firstly to avoid failures - it.increment(ec); - if (ec) - break; - - // Recursively clean - safe_remove_all(entry_path); - } - std::filesystem::remove(path, ec); -} - -} // deep_gemm diff --git a/deep_gemm/__init__.py b/deep_gemm/__init__.py index f5e52b2d12..d2cd8e76da 100644 --- a/deep_gemm/__init__.py +++ b/deep_gemm/__init__.py @@ -1,5 +1,4 @@ import os -import subprocess import torch # Set some default environment provided at setup @@ -23,64 +22,70 @@ set_block_size_multiple_of, set_pdl, get_pdl, + use_deterministic_algorithms, ) # cuBLASLt Kernels from ._C import ( cublaslt_gemm_nt, cublaslt_gemm_nn, cublaslt_gemm_tn, cublaslt_gemm_tt, + cublaslt_nvfp4_gemm_nt, + batched_syrk, batched_symm, ) -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, - # MegaMoE - get_block_m_for_mega_moe, - ) - - # 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 - pass +# DeepGEMM Kernels +from ._C import ( + # FP8-FP4 GEMMs (accept any FP8/FP4 operand combination) + 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/FP4 GEMMs (alias) + fp4_gemm_nt, + fp8_gemm_nt, fp8_gemm_nn, + fp8_gemm_tn, fp8_gemm_tt, + fp8_gemm_nt_skip_head_mid, + m_grouped_fp4_gemm_nt_contiguous, + m_grouped_fp8_gemm_nt_contiguous, + m_grouped_fp8_gemm_nn_contiguous, + m_grouped_fp4_gemm_nt_masked, + m_grouped_fp8_gemm_nt_masked, + k_grouped_fp4_gemm_nt_contiguous, + 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, + # MegaGate kernels + bf16_mega_gate, get_bf16_mega_gate_config, + # Einsum kernels + einsum, + fp8_einsum, + # Attention kernels + fp8_fp4_mqa_logits, + get_mqa_logits_metadata, + get_paged_mqa_logits_metadata, + get_sparse_mqa_logits_metadata, + get_paged_sparse_mqa_logits_metadata, + fp8_fp4_sparse_mqa_logits, + fp8_fp4_paged_sparse_mqa_logits, + fp8_fp4_paged_mqa_logits, + # Attention kernels (legacy) + fp8_mqa_logits, + fp8_paged_mqa_logits, + # Hyperconnection kernels + mega_mhc, + tf32_hc_prenorm_gemm, + # Layout kernels + transform_sf_into_required_layout, + # MegaMoE + get_block_m_for_mega_moe, +) # Mega kernels from .mega import ( @@ -103,27 +108,6 @@ 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 -) +_C.init(os.path.dirname(os.path.abspath(__file__))) -__version__ = '2.6.1' +__version__ = '2.8.0' diff --git a/deep_gemm/include/deep_gemm/comm/barrier.cuh b/deep_gemm/include/deep_gemm/comm/barrier.cuh index 50b894e937..ff683442be 100644 --- a/deep_gemm/include/deep_gemm/comm/barrier.cuh +++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh @@ -11,6 +11,18 @@ namespace deep_gemm::comm { // 60s timeout, at 2 GHz constexpr int64_t kNumTimeoutCycles = 60ll * 2000000000ll; +// Spin until `pred()` holds; on timeout, `print_timeout()` runs before the assertion +template +CUTLASS_DEVICE void wait_until(const pred_t& pred, const print_timeout_t& print_timeout) { + const auto start_clock = clock64(); + while (not pred()) { + if (clock64() - start_clock >= kNumTimeoutCycles) { + print_timeout(); + DG_DEVICE_ASSERT(false and "Timeout"); + } + } +} + CUTLASS_DEVICE void cluster_sync_with_relaxed_arrive() { // Perform cluster_sync with `barrier.cluster.arrive.relaxed` // This is slightly faster than `cute::cluster_sync` but has weaker memory ordering guarantee @@ -30,15 +42,10 @@ CUTLASS_DEVICE void grid_sync(const layout::Workspace& workspace, const auto old_value = ptx::atomic_add_rel( count_ptr, sm_idx == 0 ? (kFinishSumTag - (kNumSMs - 1)) : 1); uint32_t new_value; - const auto start_clock = clock64(); - do { - new_value = ptx::ld_acq(count_ptr); - if (clock64() - start_clock >= kNumTimeoutCycles) { - printf("DeepGEMM grid sync timeout: sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", - sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); - DG_DEVICE_ASSERT(false and "Grid sync timeout"); - } - } while (((new_value ^ old_value) & kFinishSumTag) == 0); + wait_until([&]() { return ((new_value = ptx::ld_acq(count_ptr)) ^ old_value) & kFinishSumTag; }, [&]() { + printf("DeepGEMM grid sync timeout: sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", + sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); + }); } sync_scope(); } @@ -72,14 +79,10 @@ CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, 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: 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"); - } - } + wait_until([&]() { return ptx::ld_acq_sys(signal_ptr) == target; }, [&]() { + printf("DeepGEMM NVLink barrier timeout: 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); + }); } } diff --git a/deep_gemm/include/deep_gemm/common/math.cuh b/deep_gemm/include/deep_gemm/common/math.cuh index 0f0d250481..9d8e78137d 100644 --- a/deep_gemm/include/deep_gemm/common/math.cuh +++ b/deep_gemm/include/deep_gemm/common/math.cuh @@ -1,8 +1,11 @@ #pragma once +#include #include +#include #include #include +#include namespace deep_gemm::math { @@ -70,49 +73,127 @@ CUTLASS_HOST_DEVICE float fast_rcp(const float& x) { /// Casting template -CUTLASS_DEVICE int cast_into_bf16_and_pack(old_t& x, old_t& y) { - auto bf16x2 = __float22bfloat162_rn({*reinterpret_cast(&x), *reinterpret_cast(&y)}); - return *reinterpret_cast(&bf16x2); -} - -CUTLASS_DEVICE float fast_pow2(const int& x) { - uint32_t bits_x = (x + 127) << 23; - return *reinterpret_cast(&bits_x); +CUTLASS_DEVICE nv_bfloat162 cast_into_bf16x2(old_t& x, old_t& y) { + return __float22bfloat162_rn({*reinterpret_cast(&x), *reinterpret_cast(&y)}); } -CUTLASS_DEVICE int fast_log2_ceil(float x) { - const auto bits = *reinterpret_cast(&x); - const auto exp = bits >> 23; - const auto man = bits & ((1 << 23) - 1); - return exp - 127 + (man != 0); +template +CUTLASS_DEVICE int cast_into_bf16_and_pack(old_t& x, old_t& y) { + auto bf16x2 = cast_into_bf16x2(x, y); + return *reinterpret_cast(&bf16x2); } -template -CUTLASS_DEVICE void get_e4m3_sf_and_sf_inv(const float2& amax, float2& sf, float2& sf_inv) { - DG_STATIC_ASSERT(kUseUE8M0, "Must use UE8M0"); - const float2 finfo_factor = {1.0 / 448.0, 1.0 / 448.0}; - 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); +/// Packed BF16 quantization +// Tree-reduce the amax of packed BF16 pairs (HMNMX2) +template +CUTLASS_DEVICE nv_bfloat16 get_packed_bf16_amax(const nv_bfloat162 (&values)[kNumPacked]) { + DG_STATIC_ASSERT(kNumPacked >= 2 and (kNumPacked & (kNumPacked - 1)) == 0, + "The tree requires a power-of-two size"); + nv_bfloat162 tree[kNumPacked / 2]; + #pragma unroll + for (uint32_t i = 0; i < kNumPacked / 2; ++ i) + tree[i] = __hmax2(__habs2(values[i]), __habs2(values[i + kNumPacked / 2])); + #pragma unroll + for (uint32_t stride = kNumPacked / 4; stride > 0; stride /= 2) { + #pragma unroll + for (uint32_t i = 0; i < stride; ++ i) + tree[i] = __hmax2(tree[i], tree[i + stride]); + } + return __hmax(tree[0].x, tree[0].y); +} + +// Scale two BF16 pairs by power-of-two BF16 sf_invs (exact) and pack into four FP8 bytes +CUTLASS_DEVICE uint32_t scale_bf16x2_into_fp8x4(const nv_bfloat162& lower, const nv_bfloat162& upper, + const nv_bfloat162& sf_inv_lower, const nv_bfloat162& sf_inv_upper) { + return __nv_fp8x4_e4m3(__hmul2(lower, sf_inv_lower), __hmul2(upper, sf_inv_upper)).__x; +} + +// Select a power-of-two UE8M0 SF mapping `amax` into the finite range of `quant_dtype_t`: +// the carry of the integer addition performs the exponent ceiling (carry iff the amax +// mantissa exceeds the max finite value's mantissa) +template +CUTLASS_DEVICE uint32_t get_ue8m0_sf_exp(const dtype_t& amax) { + constexpr bool kIsFP32 = cute::is_same_v; + DG_STATIC_ASSERT((kIsFP32 or cute::is_same_v), "The input type must be FP32 or BF16"); + DG_STATIC_ASSERT((cute::is_same_v or cute::is_same_v), + "The quantized type must be E4M3 or E2M1"); + constexpr bool kIsFP8 = cute::is_same_v; + constexpr uint32_t kMantissaBits = kIsFP32 ? 23 : 7; + constexpr uint32_t kMantissaMask = (1u << kMantissaBits) - 1; + constexpr uint32_t kQuantMaxMantissa = (kIsFP8 ? 0x60u : 0x40u) << (kMantissaBits - 7); // mantissa(1.75) or mantissa(1.5) + constexpr uint32_t kQuantMaxExponent = kIsFP8 ? 8 : 2; // 448 = 1.75 * 2 ^ 8, 6 = 1.5 * 2 ^ 2 + // The exponent floors match the per-token cast kernel's amax clamps: + // `max(amax, 1e-4)` for E4M3 (2 ^ -22), `max(amax, 6 * 2 ^ -126)` for E2M1 (2 ^ -126) + constexpr uint32_t kMinSFExponent = kIsFP8 ? 105 : 1; + uint32_t amax_bits; + if constexpr (kIsFP32) + amax_bits = __float_as_uint(amax); + else + amax_bits = __bfloat16_as_ushort(amax); + const auto rounded_exp = (amax_bits + kMantissaMask - kQuantMaxMantissa) >> kMantissaBits; + return cute::max(rounded_exp, kMinSFExponent + kQuantMaxExponent) - kQuantMaxExponent; +} + +// The `dtype_t` sf_inv of a UE8M0 SF exponent (reciprocal biased exponents sum to 254): +// a power of two, exact in both output formats +template +CUTLASS_DEVICE auto get_ue8m0_sf_inv(const uint32_t& sf_exp) { + if constexpr (cute::is_same_v) { + return __uint_as_float((254u - sf_exp) << 23); + } else { + DG_STATIC_ASSERT((cute::is_same_v), "The output type must be FP32 or BF16"); + return __ushort_as_bfloat16((254u - sf_exp) << 7); + } } /// Reduction -CUTLASS_DEVICE uint32_t warp_inclusive_sum(uint32_t value, const uint32_t& lane_idx) { +template +CUTLASS_DEVICE T warp_inclusive_sum(T value, const uint32_t lane_idx) { #pragma unroll for (uint32_t offset = 1; offset < 32; offset <<= 1) { - const uint32_t synced = __shfl_up_sync(0xffffffff, value, offset); + const T synced = __shfl_up_sync(0xffffffff, value, offset); if (lane_idx >= offset) value += synced; } return value; } +// Block-wide prefix sum. Synchronize the block before reusing `warp_sums`. +template +CUTLASS_DEVICE T cta_exclusive_sum(const T value, T* warp_sums, T& total) { + constexpr uint32_t kNumWarps = kNumThreads / 32; + DG_STATIC_ASSERT(kNumThreads >= 32 and kNumThreads % 32 == 0 and kNumWarps <= 32, "Invalid CTA scan shape"); + + const uint32_t lane_idx = ptx::get_lane_idx(); + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const T lane_sum = warp_inclusive_sum(value, lane_idx); + if (lane_idx == 31) + warp_sums[warp_idx] = lane_sum; + __syncthreads(); + + const T warp_total = lane_idx < kNumWarps ? warp_sums[lane_idx] : 0; + const T warp_sum = warp_inclusive_sum(warp_total, lane_idx); + total = ptx::exchange(warp_sum, kNumWarps - 1); + return lane_sum - value + ptx::exchange(warp_sum - warp_total, warp_idx); +} + +template +CUTLASS_DEVICE T cta_exclusive_sum(const T value, T* warp_sums) { + T total; + return cta_exclusive_sum(value, warp_sums, total); +} + // Operation functors +// NOTES: `fmaxf`/`fminf` lower to single `FMNMX` instructions; the generic float ternaries +// cannot (mismatched NaN semantics) and lower to `FSETP + FSEL` pairs instead, +// doubling both the instruction count and the dependency chain latency template struct ReduceSum { CUTLASS_DEVICE T operator()(T a, T b) const { return a + b; } }; template struct ReduceMax { CUTLASS_DEVICE T operator()(T a, T b) const { return a > b ? a : b; } }; template struct ReduceMin { CUTLASS_DEVICE T operator()(T a, T b) const { return a < b ? a : b; } }; +template <> struct ReduceMax { CUTLASS_DEVICE float operator()(float a, float b) const { return fmaxf(a, b); } }; +template <> struct ReduceMin { CUTLASS_DEVICE float operator()(float a, float b) const { return fminf(a, b); } }; +template <> struct ReduceMax { CUTLASS_DEVICE nv_bfloat16 operator()(nv_bfloat16 a, nv_bfloat16 b) const { return __hmax(a, b); } }; +template <> struct ReduceMax { CUTLASS_DEVICE nv_bfloat162 operator()(nv_bfloat162 a, nv_bfloat162 b) const { return __hmax2(a, b); } }; template struct ReduceAnd { CUTLASS_DEVICE T operator()(T a, T b) const { return a & b; } }; template struct ReduceOr { CUTLASS_DEVICE T operator()(T a, T b) const { return a | b; } }; @@ -144,6 +225,12 @@ template (value, ReduceSum{}); } + +template +CUTLASS_DEVICE T warp_reduce_max(T value) { + return warp_reduce(value, ReduceMax{}); +} + #endif } // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/common/packing.cuh b/deep_gemm/include/deep_gemm/common/packing.cuh new file mode 100644 index 0000000000..6f54674eb2 --- /dev/null +++ b/deep_gemm/include/deep_gemm/common/packing.cuh @@ -0,0 +1,14 @@ +#pragma once + +#include +#include + +namespace deep_gemm { + +template +constexpr CUTLASS_HOST_DEVICE uint32_t get_smem_pack_factor() { + // Packed FP4 stores two logical elements per byte in shared memory. + return cute::is_same_v ? 2 : 1; +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/common/ring_pipeline.cuh b/deep_gemm/include/deep_gemm/common/ring_pipeline.cuh new file mode 100644 index 0000000000..3836a27cf4 --- /dev/null +++ b/deep_gemm/include/deep_gemm/common/ring_pipeline.cuh @@ -0,0 +1,40 @@ +#pragma once + +#include + +#include + +#include + +#include + +namespace deep_gemm { + +template +struct RingPipeline { + DG_STATIC_ASSERT(kNumStages > 0, "Ring pipeline must contain at least one stage"); + + uint32_t stage_idx = 0; + uint32_t phase = 0; + + // Callers advance by at most one full ring, so the phase toggles at most once. + CUTLASS_DEVICE cute::tuple advance(const uint32_t step = 1) { + const uint32_t current_stage_idx = stage_idx; + const uint32_t current_phase = phase; + const uint32_t next_stage_idx = stage_idx + step; + // Modulo and division by a power of two lower to bit operations. + if constexpr ((kNumStages & (kNumStages - 1)) == 0) { + stage_idx = next_stage_idx % kNumStages; + phase ^= next_stage_idx / kNumStages; + } else { + stage_idx = next_stage_idx; + if (stage_idx >= kNumStages) { + stage_idx -= kNumStages; + phase ^= 1u; + } + } + return cute::tuple(current_stage_idx, current_phase); + } +}; + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/common/tma_copy.cuh b/deep_gemm/include/deep_gemm/common/tma_copy.cuh index 2c5bf708d4..4e3fc7c03c 100644 --- a/deep_gemm/include/deep_gemm/common/tma_copy.cuh +++ b/deep_gemm/include/deep_gemm/common/tma_copy.cuh @@ -5,52 +5,60 @@ #include #include +#include namespace deep_gemm::tma { -template +template constexpr uint32_t get_inner_block_atom_size() { - return kSwizzleMode == 0 ? BLOCK_INNER : kSwizzleMode / sizeof(dtype_t); + return kSwizzleMode == 0 ? BLOCK_INNER / kPackFactor : kSwizzleMode / sizeof(dtype_t); } template + typename dtype_t, bool kIs3DTMA = false, + uint64_t kCacheHint = static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL)> CUTLASS_DEVICE void copy(void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr, dtype_t* smem_ptr, const uint32_t& inner_idx, const uint32_t& outer_idx, const uint32_t& num_tma_multicast = 1, const uint32_t& batch_idx = 0) { DG_STATIC_ASSERT(static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL) == static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), "Invalid cache hint"); - constexpr uint32_t BLOCK_INNER_ATOM = get_inner_block_atom_size(); + constexpr uint32_t kPackFactor = get_smem_pack_factor(); + // NOTES: for sub-byte packed types, the SMEM atom stride stays in + // bytes (`sizeof(dtype_t) == 1`), while the GMEM index and loop count run over logical elements + DG_STATIC_ASSERT(kPackFactor == 1 or sizeof(dtype_t) == 1, "Packing expects a 1-byte storage type"); + constexpr uint32_t BLOCK_INNER_ATOM_STORAGE = get_inner_block_atom_size(); + constexpr uint32_t BLOCK_INNER_ATOM_LOGICAL = BLOCK_INNER_ATOM_STORAGE * kPackFactor; + DG_STATIC_ASSERT(BLOCK_INNER % BLOCK_INNER_ATOM_LOGICAL == 0, "TMA inner block must contain whole atoms"); if constexpr (not kIs3DTMA) { if (num_tma_multicast == 1) { #pragma unroll - for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) { + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++ i) { cute::SM90_TMA_LOAD_2D::copy(desc_ptr, reinterpret_cast(barrier_ptr), - static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), - smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, - inner_idx + i * BLOCK_INNER_ATOM, outer_idx); + kCacheHint, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx); } } else { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) // 2-CTA function will send signals to the leader CTA only #pragma unroll - for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) { + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++ i) { cute::SM100_TMA_2SM_LOAD_2D::copy(desc_ptr, reinterpret_cast(barrier_ptr), - static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), - smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, - inner_idx + i * BLOCK_INNER_ATOM, outer_idx); + kCacheHint, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx); } #elif (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) if (cute::block_rank_in_cluster() == 0) { #pragma unroll - for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) { + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++ i) { cute::SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, reinterpret_cast(barrier_ptr), - (1 << num_tma_multicast) - 1, static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL), - smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, - inner_idx + i * BLOCK_INNER_ATOM, outer_idx); + (1 << num_tma_multicast) - 1, kCacheHint, + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx); } } #endif @@ -58,30 +66,30 @@ copy(void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr } else { if (num_tma_multicast == 1) { #pragma unroll - for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) { + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++ i) { cute::SM90_TMA_LOAD_3D::copy(desc_ptr, reinterpret_cast(barrier_ptr), static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), - smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, - inner_idx + i * BLOCK_INNER_ATOM, outer_idx, batch_idx); + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx, batch_idx); } } else { #if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) // 2-CTA function will send signals to the leader CTA only #pragma unroll - for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) { + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++ i) { cute::SM100_TMA_2SM_LOAD_3D::copy(desc_ptr, reinterpret_cast(barrier_ptr), static_cast(cute::TMA::CacheHintSm100::EVICT_NORMAL), - smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, - inner_idx + i * BLOCK_INNER_ATOM, outer_idx, batch_idx); + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx, batch_idx); } #elif (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) if (cute::block_rank_in_cluster() == 0) { #pragma unroll - for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) { + for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM_LOGICAL; ++ i) { cute::SM90_TMA_LOAD_MULTICAST_3D::copy(desc_ptr, reinterpret_cast(barrier_ptr), (1 << num_tma_multicast) - 1, static_cast(cute::TMA::CacheHintSm90::EVICT_NORMAL), - smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, - inner_idx + i * BLOCK_INNER_ATOM, outer_idx, batch_idx); + smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM_STORAGE, + inner_idx + i * BLOCK_INNER_ATOM_LOGICAL, outer_idx, batch_idx); } } #endif diff --git a/deep_gemm/include/deep_gemm/common/types.cuh b/deep_gemm/include/deep_gemm/common/types.cuh index ee4e14c031..893a0aca89 100644 --- a/deep_gemm/include/deep_gemm/common/types.cuh +++ b/deep_gemm/include/deep_gemm/common/types.cuh @@ -7,16 +7,9 @@ namespace deep_gemm { enum class MmaKind { BF16 = 0, MXFP8FP4 = 1, + MXF4 = 2, }; -constexpr CUTLASS_HOST_DEVICE int get_element_size(const MmaKind& mma_kind) { - switch (mma_kind) { - case MmaKind::BF16: return 2; - case MmaKind::MXFP8FP4: return 1; - default: return 0; - } -} - enum class GemmType { Normal = 0, MGroupedContiguous = 1, @@ -49,4 +42,14 @@ enum class KernelType { KernelNoSF = 2 }; +// Host-visible runtime state of the `epilogue::transform` operators, which derive from +// this struct without adding members, so the host fully determines the operator on launch. +// Only the fields consumed by the selected epilogue type are meaningful +struct EpilogueArgs { + uint32_t* sfd = nullptr; + uint32_t sfd_stride = 0; + uint32_t shape_m = 0, shape_n = 0; + float alpha = 1.0f; +}; + } // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/common/utils.cuh b/deep_gemm/include/deep_gemm/common/utils.cuh index 80c12f25a9..08402888a2 100644 --- a/deep_gemm/include/deep_gemm/common/utils.cuh +++ b/deep_gemm/include/deep_gemm/common/utils.cuh @@ -21,31 +21,6 @@ struct PatternVisitor { } }; -template -CUTLASS_DEVICE void for_each_static_until(std::integer_sequence, - FuncT&& func) { - ((kIdx < kNumValid ? func.template operator()() : void()), ...); -} - -template -CUTLASS_DEVICE void for_each_static_prefix(std::integer_sequence, - const uint32_t& num_valid, - FuncT&& func) { - using seq_t = std::integer_sequence; - constexpr uint32_t kNumIndices = sizeof...(kIdx); - if constexpr (kNumIndices <= 4) { - switch (num_valid) { - case 0: break; - case 1: for_each_static_until<1>(seq_t(), func); break; - case 2: for_each_static_until<2>(seq_t(), func); break; - case 3: for_each_static_until<3>(seq_t(), func); break; - default: for_each_static_until(seq_t(), func); break; - } - } else { - ((kIdx < num_valid ? func.template operator()() : void()), ...); - } -} - template struct Vectorized { static auto zeros() { diff --git a/deep_gemm/include/deep_gemm/epilogue/clean_logits.cuh b/deep_gemm/include/deep_gemm/epilogue/clean_logits.cuh new file mode 100644 index 0000000000..9bb9c1e2fb --- /dev/null +++ b/deep_gemm/include/deep_gemm/epilogue/clean_logits.cuh @@ -0,0 +1,49 @@ +#pragma once + +#include + +#include + +namespace deep_gemm::epilogue { + +// Warp-cooperative `-inf` filler for out-of-coverage logits columns. +// The aligned interior of `[start_col, end_col)` is filled with 16-byte vectorized +// stores, while the misaligned head/tail (and too-short ranges) fall back to scalar stores. +// `kNumLanes` cooperating lanes may span multiple warps; `lane_idx` must be in `[0, kNumLanes)`. +template +struct LogitsCleaner { + static constexpr uint32_t kNumElemsPerInt4 = static_cast(sizeof(int4) / sizeof(logits_dtype_t)); + + logits_dtype_t neg_inf; + int4 neg_inf_int4; + uint32_t lane_idx; + + CUTLASS_DEVICE explicit LogitsCleaner(const uint32_t& lane_idx): lane_idx(lane_idx) { + neg_inf = -cute::numeric_limits::infinity(); + const auto neg_inf_elems = reinterpret_cast(&neg_inf_int4); + #pragma unroll + for (uint32_t i = 0; i < kNumElemsPerInt4; ++ i) + neg_inf_elems[i] = neg_inf; + } + + CUTLASS_DEVICE void fill_row(logits_dtype_t* row, const uint32_t& start_col, const uint32_t& end_col) const { + const auto aligned_start = math::align(start_col, kNumElemsPerInt4); + const auto aligned_end = math::align(end_col, kNumElemsPerInt4); + if (aligned_start >= aligned_end) { + for (uint32_t j = start_col + lane_idx; j < end_col; j += kNumLanes) + row[j] = neg_inf; + __syncwarp(); + return; + } + for (uint32_t j = start_col + lane_idx; j < aligned_start; j += kNumLanes) + row[j] = neg_inf; + for (uint32_t j = aligned_end + lane_idx; j < end_col; j += kNumLanes) + row[j] = neg_inf; + __syncwarp(); + for (uint32_t j = aligned_start + lane_idx * kNumElemsPerInt4; j < aligned_end; j += kNumLanes * kNumElemsPerInt4) + *reinterpret_cast(row + j) = neg_inf_int4; + __syncwarp(); + } +}; + +} // namespace deep_gemm::epilogue diff --git a/deep_gemm/include/deep_gemm/epilogue/sm100_mega_gate.cuh b/deep_gemm/include/deep_gemm/epilogue/sm100_mega_gate.cuh new file mode 100644 index 0000000000..7893a13d7c --- /dev/null +++ b/deep_gemm/include/deep_gemm/epilogue/sm100_mega_gate.cuh @@ -0,0 +1,282 @@ +#pragma once + +#include +#include + +#include +#include + +#include +#include +#include +#include + +namespace deep_gemm::epilogue::mega_gate { + +template CUTLASS_DEVICE float apply_scoring(const float& value) { + using ScoringType = layout::mega_gate::ScoringType; + if constexpr (kScoringType == static_cast(ScoringType::Sigmoid)) { + return 1.0f / (1.0f + expf(-value)); + } else if constexpr (kScoringType == static_cast(ScoringType::SqrtSoftplus)) { + const auto softplus_precise = log1pf(expf(value)); + const auto softplus = value > 20.0f ? value : softplus_precise; + return sqrtf(softplus); + } else { + return value; + } +} + +template +CUTLASS_DEVICE void apply_scoring(float (&values)[kNumValues]) { + using ScoringType = layout::mega_gate::ScoringType; + if constexpr (kScoringType == static_cast(ScoringType::Sigmoid)) { + float exponentials[kNumValues]; + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumValues; ++ value_idx) + exponentials[value_idx] = expf(-values[value_idx]); + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumValues; ++ value_idx) + values[value_idx] = 1.0f / (1.0f + exponentials[value_idx]); + } else if constexpr (kScoringType == static_cast(ScoringType::SqrtSoftplus)) { + float exponentials[kNumValues], softplus[kNumValues]; + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumValues; ++ value_idx) + exponentials[value_idx] = expf(values[value_idx]); + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumValues; ++ value_idx) + softplus[value_idx] = log1pf(exponentials[value_idx]); + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumValues; ++ value_idx) + values[value_idx] = sqrtf(values[value_idx] > 20.0f ? values[value_idx] : softplus[value_idx]); + } +} + +template +CUTLASS_DEVICE bool sm100_store_tmem_scores_to_global(float* scores, const uint32_t& tmem_base_addr, + const uint32_t& effective_umma_n, + const uint32_t& expert_base_idx, + const uint32_t& mma_cta_rank, + const uint32_t& gate_warp_idx) { + constexpr uint32_t kNumTmemFragmentRows = 8; + constexpr bool kPartitionedTokenPaths = UMMA_M == 128 and kNumMmaCtas == 2; + constexpr uint32_t kNumExpertsPerCta = UMMA_M / kNumMmaCtas; + const auto subpartition_idx = gate_warp_idx % 4; + const auto warp_idx_in_subpartition = gate_warp_idx / 4; + constexpr uint32_t kTokenStride = kNumGateWarps / 4 * kNumTmemFragmentRows; + const auto num_token_cols = kPartitionedTokenPaths ? effective_umma_n / 2 : effective_umma_n; + const auto token_base_idx = kPartitionedTokenPaths ? subpartition_idx / 2 * num_token_cols : 0u; + const auto expert_atom_idx = kPartitionedTokenPaths ? subpartition_idx % 2 : subpartition_idx; + const auto global_expert_idx = expert_base_idx + mma_cta_rank * kNumExpertsPerCta + + expert_atom_idx * 32 + ptx::get_lane_idx(); + const auto tmem_addr = tmem_base_addr + (subpartition_idx * 32 << 16); + + auto scores_are_finite = true; + for (auto token_idx = warp_idx_in_subpartition * kNumTmemFragmentRows; + token_idx < num_token_cols; token_idx += kTokenStride) { + uint32_t values[kNumTmemFragmentRows]; + cute::SM100_TMEM_LOAD_32dp32b8x::copy(tmem_addr + token_idx, + values[0], values[1], values[2], values[3], + values[4], values[5], values[6], values[7]); + cutlass::arch::fence_view_async_tmem_load(); + #pragma unroll + for (uint32_t row_idx = 0; row_idx < kNumTmemFragmentRows; ++ row_idx) { + const auto dst_token_idx = token_base_idx + token_idx + row_idx; + const auto score = apply_scoring(__uint_as_float(values[row_idx])); + scores_are_finite = scores_are_finite and isfinite(score); + ptx::st_global(scores + static_cast(dst_token_idx) * kNumRoutedExperts + + global_expert_idx, score); + } + } + return scores_are_finite; +} + +CUTLASS_DEVICE void warp_reduce_best(const float& score, int& expert_idx) { + const auto best_score = ptx::reduce_max_sync(score); + const auto tied_expert_idx = score == best_score and expert_idx >= 0 + ? static_cast(expert_idx) : UINT32_MAX; + expert_idx = static_cast(__reduce_min_sync(0xffffffffu, tied_expert_idx)); +} + +CUTLASS_DEVICE void compare_swap_best_first(float& lhs_score, int& lhs_expert_idx, float& rhs_score, int& rhs_expert_idx) { + if (rhs_score > lhs_score or (rhs_score == lhs_score and rhs_expert_idx < lhs_expert_idx)) { + const auto score = lhs_score; + const auto expert_idx = lhs_expert_idx; + lhs_score = rhs_score, lhs_expert_idx = rhs_expert_idx; + rhs_score = score, rhs_expert_idx = expert_idx; + } +} + +CUTLASS_DEVICE uint32_t sort4_best_first(float& score_0, float& score_1, float& score_2, float& score_3) { + int idx_0 = 0, idx_1 = 1, idx_2 = 2, idx_3 = 3; + compare_swap_best_first(score_0, idx_0, score_1, idx_1); + compare_swap_best_first(score_2, idx_2, score_3, idx_3); + compare_swap_best_first(score_0, idx_0, score_2, idx_2); + compare_swap_best_first(score_1, idx_1, score_3, idx_3); + compare_swap_best_first(score_1, idx_1, score_2, idx_2); + return static_cast(idx_0 | idx_1 << 2 | idx_2 << 4 | idx_3 << 6); +} + +template +CUTLASS_DEVICE void select_warp_topk(float (&scores)[kNumExpertWaves * 4], const int& expert_base_idx, + int& selected_expert_idx) { + uint32_t permutations[kNumExpertWaves]; + #pragma unroll + for (uint32_t expert_wave_idx = 0; expert_wave_idx < kNumExpertWaves; ++ expert_wave_idx) { + const auto local_expert_base_idx = expert_wave_idx * 4; + permutations[expert_wave_idx] = sort4_best_first(scores[local_expert_base_idx + 0], + scores[local_expert_base_idx + 1], + scores[local_expert_base_idx + 2], + scores[local_expert_base_idx + 3]); + } + + const auto lane_idx = ptx::get_lane_idx(); + uint32_t packed_cursors = 0; + #pragma unroll + for (uint32_t output_idx = 0; output_idx < kNumTopk; ++ output_idx) { + auto best_score = -cute::numeric_limits::infinity(); + int best_expert_idx = -1; + uint32_t best_expert_wave_idx = 0; + #pragma unroll + for (uint32_t expert_wave_idx = 0; expert_wave_idx < kNumExpertWaves; ++ expert_wave_idx) { + const auto cursor = packed_cursors >> (expert_wave_idx * 3) & 7; + const auto local_expert_base_idx = expert_wave_idx * 4; + const auto candidate_score = cursor == 4 ? -cute::numeric_limits::infinity() + : cursor & 2 ? cursor & 1 ? scores[local_expert_base_idx + 3] + : scores[local_expert_base_idx + 2] + : cursor & 1 ? scores[local_expert_base_idx + 1] + : scores[local_expert_base_idx + 0]; + const auto candidate_offset = permutations[expert_wave_idx] >> (cursor * 2) & 3; + const auto candidate_expert_idx = expert_base_idx + static_cast(expert_wave_idx * 128 + lane_idx * 4 + candidate_offset); + if (candidate_score > best_score) { + best_score = candidate_score; + best_expert_idx = candidate_expert_idx; + best_expert_wave_idx = expert_wave_idx; + } + } + const auto local_best_expert_idx = best_expert_idx; + warp_reduce_best(best_score, best_expert_idx); + if (local_best_expert_idx == best_expert_idx) + packed_cursors += 1u << (best_expert_wave_idx * 3); + if (lane_idx == output_idx) + selected_expert_idx = best_expert_idx; + } +} + +template +CUTLASS_DEVICE float select_warp_expert_value(const float (&values)[kNumExpertWaves * 4], + const int& expert_base_idx, const int& selected_expert_idx) { + const auto relative_expert_idx = selected_expert_idx - expert_base_idx; + const auto selected_value_idx = relative_expert_idx >= 0 + ? relative_expert_idx / 128 * 4 + relative_expert_idx % 4 : -1; + const auto source_lane_idx = static_cast(relative_expert_idx >= 0 ? relative_expert_idx % 128 / 4 : 0); + auto selected_value = 0.0f; + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumExpertWaves * 4; ++ value_idx) { + const auto exchanged = ptx::exchange(values[value_idx], source_lane_idx); + if (selected_value_idx == static_cast(value_idx)) + selected_value = exchanged; + } + return selected_value; +} + +template +CUTLASS_DEVICE void store_topk_token(const uint32_t& global_token_idx, const uint32_t& lane_idx, + const uint32_t& num_routed_experts, const uint32_t& num_shared_experts, + const uint32_t& num_duplicate_experts, const float& routed_scaling_factor, + const int& ep_rank, const int64_t& unmapped_topk_idx_stride, + const int* to_physical_map, const int* logical_count, + const int* routed_logical_count, + int64_t* topk_idx, int64_t* unmapped_topk_idx, float* topk_weights, + int selected_expert_idx, const float& selected_unbiased_score) { + DG_STATIC_ASSERT(kNumTopk <= 32, "Top-k slots exceed a warp"); + const auto num_physical_topk = kNumTopk + num_shared_experts; + if (lane_idx >= kNumTopk and lane_idx < num_physical_topk) + selected_expert_idx = static_cast(lane_idx + num_routed_experts - kNumTopk); + + auto physical_expert_idx = selected_expert_idx; + if constexpr (kToPhysicalMapExists) { + if (lane_idx < num_physical_topk) { + const auto logical_expert_idx = static_cast(selected_expert_idx); + const auto num_duplicates = static_cast((lane_idx < kNumTopk ? routed_logical_count : logical_count)[logical_expert_idx]); + DG_TRAP_ONLY_DEVICE_ASSERT(num_duplicates > 0 and num_duplicates <= num_duplicate_experts); + const auto duplicate_idx = (static_cast(ep_rank) + global_token_idx * 23333u) % num_duplicates; + physical_expert_idx = to_physical_map[logical_expert_idx * num_duplicate_experts + duplicate_idx]; + } + } + + constexpr uint32_t kSumWidth = kNumTopk <= 8 ? 8 : (kNumTopk <= 16 ? 16 : 32); + const auto topk_sum = math::warp_reduce_sum(selected_unbiased_score) + 1e-20f; + + auto selected_weight = selected_unbiased_score; + if (lane_idx < kNumTopk) { + selected_weight = selected_weight / topk_sum * routed_scaling_factor; + if constexpr (kUnmappedTopkIdxExists) + unmapped_topk_idx[global_token_idx * unmapped_topk_idx_stride + lane_idx] = selected_expert_idx; + } else if (lane_idx < num_physical_topk) { + selected_weight = 1.0f; + } + + if (lane_idx < num_physical_topk) { + const auto output_idx = global_token_idx * num_physical_topk + lane_idx; + topk_idx[output_idx] = physical_expert_idx; + topk_weights[output_idx] = selected_weight; + } +} + +CUTLASS_DEVICE uint32_t random_hash(uint32_t value) { + value ^= value >> 16; + value *= 0x7feb352du; + value ^= value >> 15; + value *= 0x846ca68bu; + return value ^ value >> 16; +} + +template +CUTLASS_DEVICE void store_random_topk_token(const uint32_t& global_token_idx, const uint32_t& lane_idx, + const uint32_t& num_routed_experts, const uint32_t& num_shared_experts, + const int* logical_count, const int& ep_rank, + const int64_t& unmapped_topk_idx_stride, int64_t* topk_idx, + int64_t* unmapped_topk_idx, float* topk_weights) { + const auto num_physical_topk = kNumTopk + num_shared_experts; + auto num_physical_experts = num_routed_experts + num_shared_experts; + if constexpr (kToPhysicalMapExists) { + auto local_count = 0u; + for (uint32_t logical_expert_idx = lane_idx; logical_expert_idx < num_physical_experts; logical_expert_idx += 32) + local_count += static_cast(logical_count[logical_expert_idx]); + num_physical_experts = __reduce_add_sync(0xffffffffu, local_count); + } + DG_TRAP_ONLY_DEVICE_ASSERT(num_physical_experts >= num_physical_topk); + + const auto seed = random_hash(static_cast(ep_rank) ^ global_token_idx * 0x9e3779b9u ^ lane_idx * 0x85ebca6bu); + const auto random_weight = __uint_as_float(0x3f800000u | random_hash(seed ^ 0xc2b2ae35u) >> 9) - 1.0f + 0x1p-24f; + constexpr uint32_t kInvalidRandomIdx = INT32_MAX; + auto candidate_idx = lane_idx < num_physical_topk + ? random_hash(seed) % (num_physical_experts - lane_idx) + : kInvalidRandomIdx; + auto selected_expert_idx = kInvalidRandomIdx; + for (uint32_t output_idx = 0; output_idx < num_physical_topk; ++ output_idx) { + const auto min_candidate_idx = __reduce_min_sync(0xffffffffu, candidate_idx); + const auto min_lane_idx = __reduce_min_sync(0xffffffffu, candidate_idx == min_candidate_idx ? lane_idx : kInvalidRandomIdx); + if (selected_expert_idx == kInvalidRandomIdx) { + if (candidate_idx >= min_candidate_idx and min_lane_idx < lane_idx) { + ++ candidate_idx; + } else if (min_lane_idx == lane_idx) { + selected_expert_idx = candidate_idx; + candidate_idx = kInvalidRandomIdx; + } + } + } + + if (lane_idx < num_physical_topk) { + const auto output_idx = global_token_idx * num_physical_topk + lane_idx; + topk_idx[output_idx] = static_cast(selected_expert_idx); + topk_weights[output_idx] = random_weight; + if constexpr (kUnmappedTopkIdxExists) { + if (lane_idx < kNumTopk) + unmapped_topk_idx[global_token_idx * unmapped_topk_idx_stride + lane_idx] = -1; + } + } +} + +} // namespace deep_gemm::epilogue::mega_gate diff --git a/deep_gemm/include/deep_gemm/epilogue/sm100_mega_mhc.cuh b/deep_gemm/include/deep_gemm/epilogue/sm100_mega_mhc.cuh new file mode 100644 index 0000000000..24960b7710 --- /dev/null +++ b/deep_gemm/include/deep_gemm/epilogue/sm100_mega_mhc.cuh @@ -0,0 +1,428 @@ +#pragma once + +#include +#include + +#include +#include +#include +#include + +namespace deep_gemm::epilogue::mega_mhc { + +using namespace layout::mega_mhc; + +// Reduce one token's split partials and produce its Mix coefficients. +template +CUTLASS_DEVICE float run_mix_task(const Workspace& workspace, uint64_t* gmem_split_barriers, + const MixArgs& mix_args, + const uint32_t& token_idx, const uint32_t& lane_idx) { + const uint32_t m_block_idx = token_idx / BLOCK_M; + const uint32_t row_idx = token_idx % BLOCK_M; + sched::mega_mhc::Mix::wait(gmem_split_barriers, m_block_idx); + + // Reduce HC norm statistics and GEMM outputs across the split tasks. + const bool is_hc_output_lane = lane_idx < kNumHCOutputs; + const uint32_t first_task_idx = m_block_idx * kNumSplits; + const auto hc_norm_sqr_sum_partials = workspace.get_hc_norm_sqr_sum_partial_ptr(first_task_idx); + float hc_norm_sqr_sum = 0.0f; + #pragma unroll + for (uint32_t k_split_idx = lane_idx; k_split_idx < kNumSplits; k_split_idx += 32) + hc_norm_sqr_sum += hc_norm_sqr_sum_partials[ + k_split_idx * kNumSqrSumPartialElementsPerTask + row_idx]; + hc_norm_sqr_sum = math::warp_reduce_sum(hc_norm_sqr_sum); + + const uint32_t safe_hc_output_idx = cute::min(lane_idx, kNumHCOutputs - 1); + const auto gemm_partials = workspace.get_gemm_partial_ptr(first_task_idx); + float hc_output = 0.0f; + #pragma unroll + for (uint32_t k_split_idx = 0; k_split_idx < kNumSplits; ++ k_split_idx) + hc_output += gemm_partials[k_split_idx * kNumGemmPartialElementsPerTask + + row_idx * kNumHCOutputs + safe_hc_output_idx]; + hc_output = is_hc_output_lane ? hc_output : 0.0f; + const float hc_rms_scale = rsqrtf(hc_norm_sqr_sum * (1.0f / static_cast(kNumRoutes * kHidden)) + mix_args.hc_norm_eps); + + // Warp lanes map to [Pre routes | Post routes | Comb routes x routes | inactive]. + constexpr uint32_t kFirstCombLane = 2 * kNumRoutes; + const bool is_pre_lane = lane_idx < kNumRoutes; + const bool is_pre_or_post_lane = lane_idx < kFirstCombLane; + const bool is_comb_lane = is_hc_output_lane and not is_pre_or_post_lane; + float mix_value = 0.0f; + if (is_hc_output_lane) { + const uint32_t mix_scale_idx = is_pre_lane ? 0u : is_pre_or_post_lane ? 1u : 2u; + const float affine_value = hc_output * hc_rms_scale * mix_args.scales[mix_scale_idx] + mix_args.bases[lane_idx]; + if (is_pre_or_post_lane) { + const float sigmoid = math::fast_rcp(1.0f + __expf(-affine_value)); + mix_value = is_pre_lane ? sigmoid + mix_args.hc_pre_eps : sigmoid * mix_args.hc_post_scale; + } else { + mix_value = affine_value; + } + } + __syncwarp(); + + const auto sinkhorn_normalize = [&](float value) { + // Comb logits occupy a row-major route x route region after Pre and Post. + value = is_comb_lane ? value : 0.0f; + const float row_max = math::warp_reduce_max(value); + value = __expf(value - row_max); + float row_sum = math::warp_reduce_sum(value); + value = value * math::fast_rcp(row_sum) + mix_args.sinkhorn_eps; + + const auto get_col_sum = [](const float& element) { + // For four routes, these XORs combine adjacent rows, then the two row pairs. + float sum = element + __shfl_xor_sync(0xffffffff, element, kNumRoutes); + sum += __shfl_xor_sync(0xffffffff, sum, 6 * kNumRoutes); + return sum; + }; + + value *= math::fast_rcp(get_col_sum(value) + mix_args.sinkhorn_eps); + #pragma unroll 1 + for (uint32_t iteration = 1; iteration < mix_args.num_sinkhorn_iters; ++ iteration) { + row_sum = math::warp_reduce_sum(value); + value *= math::fast_rcp(row_sum + mix_args.sinkhorn_eps); + value *= math::fast_rcp(get_col_sum(value) + mix_args.sinkhorn_eps); + } + return value; + }; + + const float normalized_comb = sinkhorn_normalize(mix_value); + + // Only Shifted stores the new Pre coefficients. + if (is_hc_output_lane and (kIsShifted or not is_pre_lane)) { + float* mix_output = is_pre_lane ? mix_args.new_prev_mix : + is_comb_lane ? mix_args.new_comb_res_mix : mix_args.new_post_mix; + const uint32_t mix_output_stride = is_comb_lane ? kNumRoutes * kNumRoutes : kNumRoutes; + const uint32_t mix_output_idx = is_comb_lane ? lane_idx - kFirstCombLane : lane_idx & (kNumRoutes - 1); + mix_output[token_idx * mix_output_stride + mix_output_idx] = is_comb_lane ? normalized_comb : mix_value; + } + __syncwarp(); + + return mix_value; +} + +// Every warp-wide pack spans 256 hidden elements. +constexpr uint32_t kHiddenPerNormPack = 256; +constexpr uint32_t kHiddenPerNormLane = kHiddenPerNormPack / 32; +constexpr uint32_t kNumBF16PairsPerLane = sizeof(uint4) / sizeof(nv_bfloat162); +constexpr uint32_t kNumLanesPerSF = kHiddenPerSF / kHiddenPerNormLane; +constexpr uint32_t kNumSFWordsPerNormPack = kHiddenPerNormPack / kHiddenPerSFWord; +constexpr uint32_t kNumLanesPerSFWord = kHiddenPerSFWord / kHiddenPerNormLane; + +// Shared normalization and FP8 output helpers. +CUTLASS_DEVICE nv_bfloat162 normalize_bf16_pair(const nv_bfloat162& input, const nv_bfloat162& weight, const float& scale) { + return __float22bfloat162_rn(__fmul2_rn(__fmul2_rn(__bfloat1622float2(input), {scale, scale}), __bfloat1622float2(weight))); +} + +template +CUTLASS_DEVICE uint32_t get_shared_sf_row_idx(const uint32_t& token_idx) { + // Mega MoE SF pages interleave four 32-row groups within each 128-row page. + const uint32_t idx = token_idx % SF_BLOCK_M; + return token_idx / SF_BLOCK_M * ((SF_BLOCK_M + 127u) & ~127u) + + (idx & ~127u) + (idx & 31u) * 4 + ((idx >> 5) & 3u); +} + +CUTLASS_DEVICE uint32_t store_fp8_and_pack_sf(const NormArgs& norm_args, const uint64_t& output_offset, + const uint4& values, const uint32_t& sf_exp) { + const auto bf16x2 = reinterpret_cast(&values); + // Splat the BF16 sf_inv into both packed halves + const auto sf_inv = __bfloat162bfloat162(math::get_ue8m0_sf_inv(sf_exp)); + const uint2 fp8_values = {math::scale_bf16x2_into_fp8x4(bf16x2[0], bf16x2[1], sf_inv, sf_inv), + math::scale_bf16x2_into_fp8x4(bf16x2[2], bf16x2[3], sf_inv, sf_inv)}; + *reinterpret_cast(norm_args.y_fp8 + output_offset) = fp8_values; + + // Pack four 32-element SF groups from each 16-lane group into one int32. + const uint32_t paired_sf = __byte_perm(sf_exp, __shfl_xor_sync(0xffffffff, sf_exp, kNumLanesPerSF), 0x1140); + return __byte_perm(paired_sf, __shfl_xor_sync(0xffffffff, paired_sf, 2 * kNumLanesPerSF), 0x5410); +} + +// A Normal task fuses Mix, Pre, and RMSNorm for one token. +template +CUTLASS_DEVICE void run_normal_worker(SharedStorage& smem, const Workspace& workspace, + const uint32_t& wg_idx_in_cta, + const uint32_t& warp_idx_in_wg, const uint32_t& lane_idx, + uint64_t* gmem_split_barriers, + const MixArgs& mix_args, + const NormArgs& norm_args) { + constexpr uint32_t kNumWGsPerCTA = kNumThreads / (kNumWarpsPerWG * 32); + constexpr uint32_t kNumHiddenPacksPerWarp = kHidden / (kHiddenPerNormPack * kNumWarpsPerWG); + DG_STATIC_ASSERT(kHidden % (kHiddenPerNormPack * kNumWarpsPerWG) == 0, "Hidden size must divide evenly across Normal warps"); + constexpr uint32_t kNumWarpSqrSums = kNumWGsPerCTA * kNumWarpsPerWG; + DG_STATIC_ASSERT(kNumWarpSqrSums * sizeof(float) + + kNumWGsPerCTA * (sizeof(uint32_t) + kNumRoutes * sizeof(float)) <= sizeof(smem.normal_scratch), + "Normal worker scratch is too small"); + + // Scratch is laid out as [warp sums | token indices | Pre coefficients]. + const uint32_t hidden_pack_begin = warp_idx_in_wg * kNumHiddenPacksPerWarp; + const uint32_t worker_barrier_idx = 1 + wg_idx_in_cta; + const auto warp_sqr_sums = smem.normal_scratch; + const auto token_indices = reinterpret_cast(warp_sqr_sums + kNumWarpSqrSums); + const auto pre_coeffs = reinterpret_cast(token_indices + kNumWGsPerCTA) + wg_idx_in_cta * kNumRoutes; + + while (true) { + if (warp_idx_in_wg == 0) { + uint32_t task_ticket; + if (cute::elect_one_sync()) + task_ticket = atomicAdd(&smem.next_norm_task_ticket, 1u); + task_ticket = ptx::exchange(task_ticket, 0); + const uint32_t token_idx = blockIdx.x + task_ticket * kNumSMs; + token_indices[wg_idx_in_cta] = token_idx; + + // Mix is warp-sized; only its Pre coefficients must be broadcast. + if (token_idx < norm_args.num_tokens) { + const float pre_coeff = run_mix_task( + workspace, gmem_split_barriers, mix_args, token_idx, lane_idx); + if (lane_idx < kNumRoutes) + pre_coeffs[lane_idx] = pre_coeff; + } + } + cutlass::arch::NamedBarrier::sync(kNumWarpsPerWG * 32, worker_barrier_idx); + const uint32_t token_idx = token_indices[wg_idx_in_cta]; + if (token_idx >= norm_args.num_tokens) + return; + const float pre_coeff = lane_idx < kNumRoutes ? pre_coeffs[lane_idx] : 0.0f; + + // Accumulate routes in FP32, then derive RMS from the BF16-rounded values. + const auto token_residual = norm_args.new_residual + static_cast(token_idx) * kNumRoutes * kHidden; + uint4 retained[kNumHiddenPacksPerWarp]; + float2 sqr_sum = make_float2(0.0f, 0.0f); + #pragma unroll + for (uint32_t hidden_pack_offset = 0; hidden_pack_offset < kNumHiddenPacksPerWarp; ++ hidden_pack_offset) { + const uint32_t hidden_pack_idx = hidden_pack_begin + hidden_pack_offset; + const uint32_t hidden_idx = hidden_pack_idx * kHiddenPerNormPack + lane_idx * kHiddenPerNormLane; + const auto output = reinterpret_cast(&retained[hidden_pack_offset]); + float2 values[kNumBF16PairsPerLane] = {}; + #pragma unroll + for (uint32_t route_idx = 0; route_idx < kNumRoutes; ++ route_idx) { + const float coeff = ptx::exchange(pre_coeff, route_idx); + const uint4 input = ptx::ld_evict_first( + reinterpret_cast(token_residual + route_idx * kHidden + hidden_idx)); + const auto pairs = reinterpret_cast(&input); + #pragma unroll + for (uint32_t pair_idx = 0; pair_idx < kNumBF16PairsPerLane; ++ pair_idx) + values[pair_idx] = __ffma2_rn(__bfloat1622float2(pairs[pair_idx]), {coeff, coeff}, values[pair_idx]); + } + #pragma unroll + for (uint32_t pair_idx = 0; pair_idx < kNumBF16PairsPerLane; ++ pair_idx) { + output[pair_idx] = __float22bfloat162_rn(values[pair_idx]); + values[pair_idx] = __bfloat1622float2(output[pair_idx]); + sqr_sum = __ffma2_rn(values[pair_idx], values[pair_idx], sqr_sum); + } + } + + const float warp_sqr_sum = math::warp_reduce_sum(sqr_sum.x + sqr_sum.y); + warp_sqr_sums[wg_idx_in_cta * kNumWarpsPerWG + warp_idx_in_wg] = warp_sqr_sum; + cutlass::arch::NamedBarrier::sync(kNumWarpsPerWG * 32, worker_barrier_idx); + float rms_scale = lane_idx < kNumWarpsPerWG ? warp_sqr_sums[wg_idx_in_cta * kNumWarpsPerWG + lane_idx] : 0.0f; + rms_scale = math::warp_reduce_sum(rms_scale); + rms_scale = ptx::exchange(rms_scale, 0); + constexpr float kInvHidden = 1.0f / static_cast(kHidden); + rms_scale = rsqrtf(rms_scale * kInvHidden + norm_args.eps) * norm_args.scale; + const uint64_t output_offset = static_cast(token_idx) * kHidden; + uint32_t shared_sf_row_idx = 0; + if constexpr (SF_BLOCK_M > 0) + shared_sf_row_idx = get_shared_sf_row_idx(token_idx); + + const auto weight_packs = reinterpret_cast( + norm_args.weight + hidden_pack_begin * kHiddenPerNormPack + lane_idx * kHiddenPerNormLane); + uint4 weight = weight_packs[0]; + #pragma unroll + for (uint32_t hidden_pack_offset = 0; hidden_pack_offset < kNumHiddenPacksPerWarp; ++ hidden_pack_offset) { + const uint32_t hidden_pack_idx = hidden_pack_begin + hidden_pack_offset; + const bool has_next_weight = hidden_pack_offset + 1 < kNumHiddenPacksPerWarp; + uint4 next_weight; + if (has_next_weight) + next_weight = weight_packs[(hidden_pack_offset + 1) * (kHiddenPerNormPack / kHiddenPerNormLane)]; + + const uint32_t hidden_idx = hidden_pack_idx * kHiddenPerNormPack + lane_idx * kHiddenPerNormLane; + const auto values = reinterpret_cast(&retained[hidden_pack_offset]); + const auto weight_pairs = reinterpret_cast(&weight); + #pragma unroll + for (uint32_t pair_idx = 0; pair_idx < kNumBF16PairsPerLane; ++ pair_idx) + values[pair_idx] = normalize_bf16_pair(values[pair_idx], weight_pairs[pair_idx], rms_scale); + if constexpr (kStoreBF16) + *reinterpret_cast(norm_args.y_bf16 + output_offset + hidden_idx) = retained[hidden_pack_offset]; + if constexpr (kStoreFP8) { + const auto amax = math::warp_reduce_max( + math::get_packed_bf16_amax(reinterpret_cast(retained[hidden_pack_offset]))); + const uint32_t sf_word_idx = hidden_pack_idx * kNumSFWordsPerNormPack + lane_idx / kNumLanesPerSFWord; + const uint32_t packed_sf = store_fp8_and_pack_sf( + norm_args, output_offset + hidden_idx, retained[hidden_pack_offset], + math::get_ue8m0_sf_exp(amax)); + if constexpr (SF_BLOCK_M > 0) { + if (lane_idx % kNumLanesPerSFWord == 0) { + norm_args.y_primary_sf[token_idx * norm_args.y_primary_sf_stride_token + sf_word_idx * norm_args.y_primary_sf_stride_word] = packed_sf; + norm_args.y_shared_sf[shared_sf_row_idx + sf_word_idx * norm_args.y_shared_sf_stride_word] = packed_sf; + } + } else { + retained[hidden_pack_offset].x = packed_sf; + } + } + + if (has_next_weight) + weight = next_weight; + } + + // Keep column-major SF in dead value registers until contiguous output stores finish. + if constexpr (kStoreFP8 and SF_BLOCK_M == 0) { + if (lane_idx % kNumLanesPerSFWord == 0) { + const int64_t sf_word_stride = norm_args.y_primary_sf_stride_word; + auto sf_ptr = norm_args.y_primary_sf + token_idx * norm_args.y_primary_sf_stride_token + + (hidden_pack_begin * kNumSFWordsPerNormPack + lane_idx / kNumLanesPerSFWord) * sf_word_stride; + + #pragma unroll + for (uint32_t hidden_pack_offset = 0; hidden_pack_offset < kNumHiddenPacksPerWarp; ++ hidden_pack_offset) { + *sf_ptr = retained[hidden_pack_offset].x; + sf_ptr += kNumSFWordsPerNormPack * sf_word_stride; + } + } + } + } +} + +// Shifted Norm consumes materialized X1 and handles two tokens per task. +constexpr uint32_t kNumTokensPerShiftedNormTask = 2; + +struct ShiftedNormPack { + uint4 x1[kNumTokensPerShiftedNormTask]; + uint4 weight; +}; + +template +CUTLASS_DEVICE void run_shifted_norm_task(const Workspace& workspace, uint64_t* gmem_split_barriers, + const NormArgs& norm_args, + const uint32_t& task_idx, + const uint32_t& num_partitions, const uint32_t& lane_idx) { + constexpr uint32_t kNumHiddenPacks = kHidden / kHiddenPerNormPack; + // task_idx flattens [token pair, hidden-pack partition]. + const uint32_t partition_idx = task_idx % num_partitions; + const uint32_t hidden_pack_begin = partition_idx * kNumHiddenPacks / num_partitions; + const uint32_t hidden_pack_end = (partition_idx + 1) * kNumHiddenPacks / num_partitions; + const uint32_t m_idx = task_idx / num_partitions * kNumTokensPerShiftedNormTask; + const uint32_t m_block_idx = m_idx / BLOCK_M; + const bool has_second_token = m_idx + 1 < norm_args.num_tokens; + uint2 shared_sf_row_indices = {}; + + sched::mega_mhc::Norm::wait(gmem_split_barriers, m_block_idx); + const auto load_pack = [&](const uint64_t output_offset, const uint32_t hidden_idx) { + return ShiftedNormPack{ + {ptx::ld_evict_first(reinterpret_cast(norm_args.y_bf16 + output_offset)), + ptx::ld_evict_first(reinterpret_cast( + norm_args.y_bf16 + output_offset + (has_second_token ? kHidden : 0)))}, + *reinterpret_cast(norm_args.weight + hidden_idx)}; + }; + + // Keep current/next packs in registers and prefetch two ahead to hide the BF16 reread. + ShiftedNormPack pack, next_pack; + const uint32_t first_hidden_idx = hidden_pack_begin * kHiddenPerNormPack + lane_idx * kHiddenPerNormLane; + const uint64_t first_output_offset = static_cast(m_idx) * kHidden + first_hidden_idx; + pack = load_pack(first_output_offset, first_hidden_idx); + if (hidden_pack_begin + 1 < hidden_pack_end) + next_pack = load_pack(first_output_offset + kHiddenPerNormPack, first_hidden_idx + kHiddenPerNormPack); + + float x1_sqr_sum = 0.0f; + if (lane_idx < kNumTokensPerShiftedNormTask and m_idx + lane_idx < norm_args.num_tokens) { + const auto x1_sqr_sum_partials = workspace.get_x1_sqr_sum_partial_ptr(m_block_idx * kNumSplits); + #pragma unroll + for (uint32_t k_split_idx = 0; k_split_idx < kNumSplits; ++ k_split_idx) + x1_sqr_sum += x1_sqr_sum_partials[k_split_idx * kNumSqrSumPartialElementsPerTask + + m_idx % BLOCK_M + lane_idx]; + } + constexpr float kInvHidden = 1.0f / static_cast(kHidden); + const float rms_scale = rsqrtf(x1_sqr_sum * kInvHidden + norm_args.eps) * norm_args.scale; + const float2 rms_scales = {ptx::exchange(rms_scale, 0), ptx::exchange(rms_scale, 1)}; + if constexpr (SF_BLOCK_M > 0) { + shared_sf_row_indices.x = get_shared_sf_row_idx(m_idx); + shared_sf_row_indices.y = get_shared_sf_row_idx(m_idx + 1); + } + for (uint32_t hidden_pack_idx = hidden_pack_begin; hidden_pack_idx < hidden_pack_end; ++ hidden_pack_idx) { + const uint32_t hidden_idx = hidden_pack_idx * kHiddenPerNormPack + lane_idx * kHiddenPerNormLane; + const uint64_t output_offset = static_cast(m_idx) * kHidden + hidden_idx; + const bool has_prefetch = hidden_pack_idx + 2 < hidden_pack_end; + ShiftedNormPack prefetch_pack; + if (has_prefetch) + prefetch_pack = load_pack(output_offset + 2 * kHiddenPerNormPack, hidden_idx + 2 * kHiddenPerNormPack); + uint4 values[kNumTokensPerShiftedNormTask] = {}; + const auto input = reinterpret_cast(pack.x1); + const auto weight = reinterpret_cast(&pack.weight); + auto output = reinterpret_cast(values); + #pragma unroll + for (uint32_t pair_idx = 0; pair_idx < kNumBF16PairsPerLane; ++ pair_idx) { + output[pair_idx] = normalize_bf16_pair(input[pair_idx], weight[pair_idx], rms_scales.x); + output[kNumBF16PairsPerLane + pair_idx] = + normalize_bf16_pair(input[kNumBF16PairsPerLane + pair_idx], weight[pair_idx], rms_scales.y); + } + if constexpr (kStoreBF16) { + *reinterpret_cast(norm_args.y_bf16 + output_offset) = values[0]; + if (has_second_token) + *reinterpret_cast(norm_args.y_bf16 + output_offset + kHidden) = values[1]; + } + if constexpr (kStoreFP8) { + const auto values_pairs = reinterpret_cast(values); + // Pack both tokens' amaxes into one BF16 pair to share a single warp reduction + const auto packed_amax = math::warp_reduce_max(nv_bfloat162( + math::get_packed_bf16_amax(values_pairs[0]), + has_second_token ? math::get_packed_bf16_amax(values_pairs[1]) : __ushort_as_bfloat16(0))); + const uint2 sf_exponents = {math::get_ue8m0_sf_exp(packed_amax.x), + math::get_ue8m0_sf_exp(packed_amax.y)}; + const uint32_t sf_word_idx = hidden_pack_idx * kNumSFWordsPerNormPack + lane_idx / kNumLanesPerSFWord; + const uint32_t first_packed_sf = store_fp8_and_pack_sf(norm_args, output_offset, values[0], sf_exponents.x); + if (lane_idx % kNumLanesPerSFWord == 0) { + norm_args.y_primary_sf[m_idx * norm_args.y_primary_sf_stride_token + sf_word_idx * norm_args.y_primary_sf_stride_word] = first_packed_sf; + if constexpr (SF_BLOCK_M > 0) + norm_args.y_shared_sf[static_cast(sf_word_idx) * norm_args.y_shared_sf_stride_word + shared_sf_row_indices.x] = first_packed_sf; + } + if (has_second_token) { + const uint32_t second_packed_sf = store_fp8_and_pack_sf( + norm_args, output_offset + kHidden, values[1], sf_exponents.y); + if (lane_idx % kNumLanesPerSFWord == 0) { + norm_args.y_primary_sf[(m_idx + 1) * norm_args.y_primary_sf_stride_token + sf_word_idx * norm_args.y_primary_sf_stride_word] = second_packed_sf; + if constexpr (SF_BLOCK_M > 0) + norm_args.y_shared_sf[static_cast(sf_word_idx) * norm_args.y_shared_sf_stride_word + shared_sf_row_indices.y] = second_packed_sf; + } + } + } + if (hidden_pack_idx + 1 < hidden_pack_end) + pack = next_pack; + if (has_prefetch) + next_pack = prefetch_pack; + } +} + +template +CUTLASS_DEVICE void run_shifted_norm_worker(SharedStorage& smem, const Workspace& workspace, + uint64_t* gmem_split_barriers, + const NormArgs& norm_args, + const uint32_t& lane_idx) { + constexpr uint32_t kNumTasksPerGridWave = kNumSMs * kNumWarpsPerWG; + constexpr uint32_t kNumHiddenPacks = kHidden / kHiddenPerNormPack; + const uint32_t num_token_pairs = math::ceil_div(norm_args.num_tokens, kNumTokensPerShiftedNormTask); + // Only subdivide underfilled FP8 grids when a token pair can span multiple CTAs. + constexpr uint32_t kNumFP8Partitions = + kNumSplits > kDefaultNumSplits and kNumHiddenPacks > 2 * kNumWarpsPerWG ? kNumHiddenPacks / 2 : 2; + const uint32_t num_partitions = kStoreFP8 ? kNumFP8Partitions : cute::max( + math::ceil_div(kNumHiddenPacks, cute::max( + 1u, num_token_pairs * kNumHiddenPacks / kNumTasksPerGridWave)), 2u); + const uint32_t num_tasks = num_token_pairs * num_partitions; + const uint32_t cta_task_base = blockIdx.x * kNumWarpsPerWG; + if (cta_task_base >= num_tasks) + return; + + cutlass::arch::warpgroup_reg_alloc(); + while (true) { + uint32_t task_ticket; + if (cute::elect_one_sync()) + task_ticket = atomicAdd(&smem.next_norm_task_ticket, 1u); + task_ticket = ptx::exchange(task_ticket, 0); + // Preserve the four CTA-local warp residue classes, then advance by one grid-wide worker wave. + const uint32_t task_idx = cta_task_base + task_ticket % kNumWarpsPerWG + + task_ticket / kNumWarpsPerWG * kNumTasksPerGridWave; + if (task_idx >= num_tasks) + return; + run_shifted_norm_task( + workspace, gmem_split_barriers, norm_args, task_idx, num_partitions, lane_idx); + } +} + +} // namespace deep_gemm::epilogue::mega_mhc diff --git a/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd.cuh b/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd.cuh index bf0e460c8f..75bef067a4 100644 --- a/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd.cuh +++ b/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd.cuh @@ -4,7 +4,7 @@ #include #include -#include +#include #include #include @@ -15,17 +15,25 @@ template + typename epilogue_op_t, + typename smem_t> CUTLASS_DEVICE void -sm100_store_cd(const utils::PatternVisitor& smem_cd, uint32_t& tma_stage_idx, +sm100_store_cd(smem_t& smem, uint32_t& tma_stage_idx, const uint32_t& tmem_base_addr, const uint32_t& base_m_idx, const uint32_t& base_n_idx, const uint32_t& batch_idx, + const bool& is_empty_group, const uint32_t& epilogue_warp_idx, const uint32_t& lane_idx, + const epilogue_op_t& epilogue_op, + const bool& reverse_store_order, + const cutlass::arch::ClusterTransactionBarrier* tmem_overlap_barrier, const cutlass::arch::ClusterTransactionBarrier* tmem_empty_barrier, const cute::TmaDescriptor& tensor_map_cd) { + using cd_dtype_t = typename smem_t::cd_dtype; + // Whether to cast D into FP8 with dynamic per-`kSFGranN` UE8M0 SFs (see `EpilogueDynamicScaledFP8`) + constexpr bool kWithOutputSF = cute::is_same_v; + constexpr uint32_t kSFGranN = transform::EpilogueDynamicScaledFP8::kSFGranN; // TMA checks constexpr uint32_t kNumBankGroupBytes = 16; constexpr uint32_t kNumElemsPerBankGroup = kNumBankGroupBytes / sizeof(cd_dtype_t); @@ -33,6 +41,8 @@ sm100_store_cd(const utils::PatternVisitor& smem_cd, uint32_t& tma DG_STATIC_ASSERT(STORE_BLOCK_N % kNumElemsPerBankGroup == 0, "Invalid swizzling"); DG_STATIC_ASSERT(BLOCK_M % STORE_BLOCK_M == 0, "Invalid block sizes"); DG_STATIC_ASSERT(BLOCK_N % STORE_BLOCK_N == 0, "Invalid block sizes"); + DG_STATIC_ASSERT(not kWithOutputSF or cute::is_same_v, "FP8 output requires an E4M3 D"); + DG_STATIC_ASSERT(not kWithOutputSF or STORE_BLOCK_N % kSFGranN == 0, "A store must cover complete SF groups"); // Share store pipeline between blocks auto advance_store_pipeline = [&]() { @@ -47,7 +57,24 @@ sm100_store_cd(const utils::PatternVisitor& smem_cd, uint32_t& tma constexpr uint32_t kNumStores = BLOCK_N / STORE_BLOCK_N; #pragma unroll for (uint32_t s = 0; s < kNumStores; ++ s, advance_store_pipeline()) { - auto smem_base_ptr = reinterpret_cast(smem_cd[tma_stage_idx]); + const auto store_idx = reverse_store_order ? kNumStores - 1 - s : s; + auto smem_base_ptr = reinterpret_cast(smem.cd[tma_stage_idx]); + + // Swizzled shared memory address of the `bank_group_idx`-th bank group in this + // warp's atom, reshaping the atom in another view: + // - original: `(LAYOUT_AD_M, kSwizzleCDMode / kNumBankGroupBytes)` + // - new: `(LAYOUT_AD_M * kSwizzleCDMode / kNumBankGroupBytes / 8, 8)` + // NOTES: "8" is the number of bank groups, "16" is the swizzling pattern + const auto get_swizzled_smem_ptr = [&](const uint32_t& bank_group_idx) { + constexpr bool kHasShortcut = (kSwizzleCDMode / kNumBankGroupBytes) == 8; + const auto shifted_idx = bank_group_idx + lane_idx * (kSwizzleCDMode / kNumBankGroupBytes); + auto row = kHasShortcut ? (bank_group_idx / 8 + lane_idx) : (shifted_idx / 8); + auto col = kHasShortcut ? bank_group_idx : (shifted_idx % 8); + col ^= row % (kSwizzleCDMode / 16); + return smem_base_ptr + // Base pointer + epilogue_warp_idx * 32 * kSwizzleCDMode + // Warp offset + row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset + }; // Wait shared memory to be released if (epilogue_warp_idx == 0) @@ -56,49 +83,95 @@ sm100_store_cd(const utils::PatternVisitor& smem_cd, uint32_t& tma // The pipeline stage const auto m_idx = base_m_idx + w * STORE_BLOCK_M; - const auto n_idx = epilogue_type_t::apply_index_n(base_n_idx + s * STORE_BLOCK_N); + const auto n_idx = epilogue_op_t::apply_index_n(base_n_idx + store_idx * STORE_BLOCK_N); + + // Iterate over TMEM loads: each covers one swizzled bank group per lane, or one + // complete SF group per lane (one row, `kSFGranN` columns) with FP8 output + constexpr uint32_t kNumElemsPerLoad = kWithOutputSF ? kSFGranN : kNumElemsPerBankGroup; + constexpr uint32_t kNumLoads = STORE_BLOCK_N / kNumElemsPerLoad; + constexpr auto kNumOverlapLoads = math::constexpr_ceil_div(kNumOverlappedTmemCols, kNumElemsPerLoad); + DG_STATIC_ASSERT(kNumOverlapLoads <= kNumLoads, "The first store block must cover all overlapped columns"); + + // A store covering a whole packed word (with word-aligned batches) writes all 4 SF bytes at once + const auto row_idx = m_idx + epilogue_warp_idx * 32 + lane_idx; + const bool store_whole_sf_word = kWithOutputSF and kNumLoads == 4 and epilogue_op.shape_n % (4 * kSFGranN) == 0; + uint32_t sf_word = 0; - // Store into shared memory #pragma unroll - for (uint32_t i = 0; i < STORE_BLOCK_N / kNumElemsPerBankGroup; ++ i) { - // Calculate the index of the bank group to be written in the atom - auto bank_group_index = i + lane_idx * (kSwizzleCDMode / kNumBankGroupBytes); - - // Reshape the atom in another view and swizzle - // - original: `(LAYOUT_AD_M, kSwizzleCDMode / kNumBankGroupBytes)` - // - new: `(LAYOUT_AD_M * kSwizzleCDMode / kNumBankGroupBytes / 8, 8)` - // NOTES: "8" is the number of bank groups, "16" is the swizzling pattern - constexpr bool kHasShortcut = (kSwizzleCDMode / kNumBankGroupBytes) == 8; - auto row = kHasShortcut ? (i / 8 + lane_idx) : (bank_group_index / 8); - auto col = kHasShortcut ? (i) : (bank_group_index % 8); - col ^= row % (kSwizzleCDMode / 16); + for (uint32_t i = 0; i < kNumLoads; ++ i) { + const auto load_idx = reverse_store_order ? kNumLoads - 1 - i : i; + const auto tmem_addr = tmem_base_addr + // Accumulator offset + w * BLOCK_N + // Wave offset + store_idx * STORE_BLOCK_N + load_idx * kNumElemsPerLoad; // In-block offset + + // Read from tensor memory into registers + // NOTES: empty groups also read the accumulator (safe, as at least one UMMA is + // issued) and select zeros afterwards, keeping the loads branchless + uint32_t values[kNumElemsPerLoad]; + DG_STATIC_ASSERT(kNumElemsPerLoad == 4 or kNumElemsPerLoad == 8 or kNumElemsPerLoad == 32, "Invalid load width"); + using tmem_load_t = cute::conditional_t>; + [&](cute::index_sequence) { + tmem_load_t::copy(tmem_addr, values[Is]...); + }(cute::make_index_sequence{}); + cutlass::arch::fence_view_async_tmem_load(); + epilogue_op.apply_values(values); + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumElemsPerLoad; ++ value_idx) + values[value_idx] = is_empty_group ? 0u : values[value_idx]; - // Source and destination memory address - uint32_t tmem_addr = tmem_base_addr + // Accumulator offset - w * BLOCK_N + // Wave offset - s * STORE_BLOCK_N + i * kNumElemsPerBankGroup; // In-block offset - auto smem_ptr = smem_base_ptr + // Base pointer - epilogue_warp_idx * 32 * kSwizzleCDMode + // Warp offset - row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset - - // Load from tensor memory, store into shared memory - uint32_t values[kNumElemsPerBankGroup]; - if constexpr (cute::is_same_v) { - // For FP32 output, read and store - DG_STATIC_ASSERT(kNumElemsPerBankGroup == 4, "Invalid type"); - cute::SM100_TMEM_LOAD_32dp32b4x::copy(tmem_addr, - values[0], values[1], values[2], values[3]); - cutlass::arch::fence_view_async_tmem_load(); - ptx::st_shared(smem_ptr, values[0], values[1], values[2], values[3]); + // Notify the MMA warp once all overlapped TMEM columns have been read, + // before any of the cast/store work + if constexpr (kNumOverlapLoads > 0) { + if (w == 0 and s == 0 and i + 1 == kNumOverlapLoads) { + ptx::tcgen05_before_thread_sync(); + tmem_overlap_barrier->arrive(0u); + } + } + // Notify tensor memory empty (only at the leader CTA) as soon as all loads are in registers + if (w == kNumMWaves - 1 and s == kNumStores - 1 and i == kNumLoads - 1) { + ptx::tcgen05_before_thread_sync(); + tmem_empty_barrier->arrive(0u); + } + + // Store into shared memory + if constexpr (kWithOutputSF) { + // Round the accumulator into BF16 pairs first, so the output bitwise matches + // a BF16 D followed by the standalone cast (power-of-two scaling of BF16 is exact) + nv_bfloat162 values_bf16x2[kNumElemsPerLoad / 2]; + #pragma unroll + for (uint32_t pair_idx = 0; pair_idx < kNumElemsPerLoad / 2; ++ pair_idx) + values_bf16x2[pair_idx] = math::cast_into_bf16x2(values[pair_idx * 2], values[pair_idx * 2 + 1]); + + // Reduce the group amax and compute the SF (the SF exponent floor covers the amax clamp), + // then splat the BF16 sf_inv into both packed halves + const auto sf_exp = math::get_ue8m0_sf_exp(math::get_packed_bf16_amax(values_bf16x2)); + const auto sf_inv = __bfloat162bfloat162(math::get_ue8m0_sf_inv(sf_exp)); + + // Cast into FP8 and store into shared memory + constexpr uint32_t kNumBankGroupsPerLoad = kNumElemsPerLoad / kNumElemsPerBankGroup; + #pragma unroll + for (uint32_t bank_group_iter_idx = 0; bank_group_iter_idx < kNumBankGroupsPerLoad; ++ bank_group_iter_idx) { + const auto smem_ptr = get_swizzled_smem_ptr(load_idx * kNumBankGroupsPerLoad + bank_group_iter_idx); + const auto pair_idx = bank_group_iter_idx * (kNumElemsPerBankGroup / 2); + ptx::st_shared(smem_ptr, + math::scale_bf16x2_into_fp8x4(values_bf16x2[pair_idx + 0], values_bf16x2[pair_idx + 1], sf_inv, sf_inv), + math::scale_bf16x2_into_fp8x4(values_bf16x2[pair_idx + 2], values_bf16x2[pair_idx + 3], sf_inv, sf_inv), + math::scale_bf16x2_into_fp8x4(values_bf16x2[pair_idx + 4], values_bf16x2[pair_idx + 5], sf_inv, sf_inv), + math::scale_bf16x2_into_fp8x4(values_bf16x2[pair_idx + 6], values_bf16x2[pair_idx + 7], sf_inv, sf_inv)); + } + + // Accumulate the SF byte into the packed word, or store it directly + if (store_whole_sf_word) + sf_word |= sf_exp << (load_idx * 8); + else + epilogue_op.store_sf(row_idx, n_idx + load_idx * kSFGranN, batch_idx, static_cast(sf_exp)); + } else if constexpr (cute::is_same_v) { + ptx::st_shared(get_swizzled_smem_ptr(load_idx), values[0], values[1], values[2], values[3]); } else { - // For BF16 output, read, cast and store - DG_STATIC_ASSERT(kNumElemsPerBankGroup == 8 and cute::is_same_v, "Invalid type"); - cute::SM100_TMEM_LOAD_32dp32b8x::copy(tmem_addr, - values[0], values[1], values[2], values[3], - values[4], values[5], values[6], values[7]); - cutlass::arch::fence_view_async_tmem_load(); ptx::st_shared( - smem_ptr, + get_swizzled_smem_ptr(load_idx), math::cast_into_bf16_and_pack(values[0], values[1]), math::cast_into_bf16_and_pack(values[2], values[3]), math::cast_into_bf16_and_pack(values[4], values[5]), @@ -106,19 +179,16 @@ sm100_store_cd(const utils::PatternVisitor& smem_cd, uint32_t& tma ); } } - - // Notify tensor memory empty (only at the leader CTA) arrival ASAP - // NOTES: only the last stage needs to do this - if (w == kNumMWaves - 1 and s == BLOCK_N / STORE_BLOCK_N - 1) { - ptx::tcgen05_before_thread_sync(); - tmem_empty_barrier->arrive(0u); + if constexpr (kWithOutputSF) { + if (store_whole_sf_word) + epilogue_op.store_sf(row_idx, n_idx, batch_idx, sf_word); } // Synchronize all threads and issue TMA cute::tma_store_fence(); cutlass::arch::NamedBarrier::sync(kNumUMMAStoreThreads, 0); if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { - if constexpr (kGemmType == GemmType::Batched) { + if constexpr (kGemmType == GemmType::Batched or is_k_grouped_contiguous(kGemmType)) { using cute_tma_t = cute::conditional_t; cute_tma_t::copy(&tensor_map_cd, smem_base_ptr, n_idx, m_idx, batch_idx); diff --git a/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh b/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh index f3f5351e6a..25711cdfbe 100644 --- a/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh +++ b/deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh @@ -4,7 +4,7 @@ #include #include -#include +#include #include #include @@ -15,22 +15,31 @@ template + typename epilogue_op_t, + typename smem_t> CUTLASS_DEVICE void -sm100_store_cd_swap_ab(const utils::PatternVisitor& smem_cd, uint32_t& tma_stage_idx, +sm100_store_cd_swap_ab(smem_t& smem, uint32_t& tma_stage_idx, const uint32_t& tmem_base_addr, const uint32_t& base_m_idx, const uint32_t& base_n_idx, const uint32_t& batch_idx, + const bool& is_empty_group, const uint32_t& effective_m, const uint32_t& epilogue_warp_idx, const uint32_t& lane_idx, + const epilogue_op_t& epilogue_op, + const bool& reverse_store_order, + const cutlass::arch::ClusterTransactionBarrier* tmem_overlap_barrier, const cutlass::arch::ClusterTransactionBarrier* tmem_empty_barrier, const cute::TmaDescriptor& tensor_map_cd) { + using cd_dtype_t = typename smem_t::cd_dtype; // NOTES: The epilogue requires a full warpgroup to read all 128 TMEM rows, // implying STORE_BLOCK_N must be 128. DG_STATIC_ASSERT(STORE_BLOCK_N == 128, "STORE_BLOCK_N must be 128 to match TMEM rows"); + // Whether to cast D into FP8 with dynamic per-`kSFGranN` UE8M0 SFs (see `EpilogueDynamicScaledFP8`) + constexpr bool kWithOutputSF = cute::is_same_v; + constexpr uint32_t kSFGranN = transform::EpilogueDynamicScaledFP8::kSFGranN; + // TMA checks constexpr uint32_t STORE_BLOCK_N_ATOM = kSwizzleCDMode / sizeof(cd_dtype_t); constexpr uint32_t kNumBankGroupBytes = 16; @@ -40,39 +49,124 @@ sm100_store_cd_swap_ab(const utils::PatternVisitor& smem_cd, uint3 DG_STATIC_ASSERT(BLOCK_N % STORE_BLOCK_N == 0, "Invalid block sizes"); DG_STATIC_ASSERT(STORE_BLOCK_M % kNumSwizzleAtomRows == 0, "Invalid swizzling"); DG_STATIC_ASSERT(STORE_BLOCK_N % STORE_BLOCK_N_ATOM == 0, "Invalid swizzling"); + DG_STATIC_ASSERT(not kWithOutputSF or cute::is_same_v, "FP8 output requires an E4M3 D"); + DG_STATIC_ASSERT(not kWithOutputSF or (STORE_BLOCK_M == 16 and kNumUMMAStoreThreads == 128), + "Swap-AB FP8 output requires one warpgroup and a 16x128 store shape"); // Share store pipeline between blocks auto advance_store_pipeline = [&]() { tma_stage_idx = (tma_stage_idx + 1) % kNumTMAStoreStages; }; - // Iterate over M blocks + // Iterate over M blocks. The scheduler skips empty blocks and aligns dynamic effective M to STORE_BLOCK_M. const auto num_stores = effective_m / STORE_BLOCK_M; for (uint32_t s = 0; s < num_stores; ++ s, advance_store_pipeline()) { + const auto store_idx = reverse_store_order ? num_stores - 1 - s : s; // Wait shared memory to be released if (epilogue_warp_idx == 0) cute::tma_store_wait(); cutlass::arch::NamedBarrier::sync(kNumUMMAStoreThreads, 0); // Store into shared memory + constexpr uint32_t kNumTmemLoads = STORE_BLOCK_M / kNumSwizzleAtomRows; #pragma unroll - for (uint32_t i = 0; i < STORE_BLOCK_M / kNumSwizzleAtomRows; ++ i) { + for (uint32_t i = 0; i < kNumTmemLoads; ++ i) { + const auto load_idx = reverse_store_order ? kNumTmemLoads - 1 - i : i; uint32_t tmem_addr = tmem_base_addr + - s * STORE_BLOCK_M + // Store stage offset - i * kNumSwizzleAtomRows; // In-block offset + store_idx * STORE_BLOCK_M + // Store stage offset + load_idx * kNumSwizzleAtomRows; // In-block offset uint32_t values[kNumSwizzleAtomRows]; // Warps cooperatively write an atomic block to shared memory DG_STATIC_ASSERT(STORE_BLOCK_N_ATOM % 32 == 0, "Invalid block sizes"); constexpr uint32_t kNumWarpsPerAtom = STORE_BLOCK_N_ATOM / 32; uint32_t outer_atom_offset = (epilogue_warp_idx / kNumWarpsPerAtom) * STORE_BLOCK_M * kSwizzleCDMode; - uint32_t inner_atom_offset = i * kNumSwizzleAtomRows * kSwizzleCDMode; - auto smem_base_ptr = reinterpret_cast(smem_cd[tma_stage_idx]) + outer_atom_offset + inner_atom_offset; + uint32_t inner_atom_offset = load_idx * kNumSwizzleAtomRows * kSwizzleCDMode; + auto smem_base_ptr = reinterpret_cast(smem.cd[tma_stage_idx]) + outer_atom_offset + inner_atom_offset; + // Read from tensor memory into registers + // NOTES: empty groups also read the accumulator (safe, as at least one UMMA is + // issued) and select zeros afterwards, keeping the loads branchless if constexpr (cute::is_same_v) { - // NOTES: Swizzling is not required in this case, but used here for consistency with other cases + // NOTES: the FP32 store does not use STSM, so the plain `.32x32b` layout works cute::SM100_TMEM_LOAD_32dp32b8x::copy(tmem_addr, values[0], values[1], values[2], values[3], values[4], values[5], values[6], values[7]); + cutlass::arch::fence_view_async_tmem_load(); + epilogue_op.apply_values(values); + } else { + // Load from TMEM using `.16x256b` shape to satisfy the STSM layout requirements + // (`.b16` for the BF16 store, `.b8` for the FP8 store): each lane receives + // 2 rows (TMEM columns `2 * (lane % 4)` onwards) x 2 columns (dps `lane / 4` and `lane / 4 + 8`) + // Start from lane index 0 + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, + values[0], values[1], values[2], values[3]); + // Start from lane index 16 + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, + values[4], values[5], values[6], values[7]); + cutlass::arch::fence_view_async_tmem_load(); + epilogue_op.apply_values(values); + } + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumSwizzleAtomRows; ++ value_idx) + values[value_idx] = is_empty_group ? 0u : values[value_idx]; + + // Notify the MMA warp once all overlapped TMEM columns have been read, + // before any of the cast/store work + constexpr auto kNumOverlapLoads = math::constexpr_ceil_div(kNumOverlappedTmemCols, kNumSwizzleAtomRows); + if constexpr (kNumOverlapLoads > 0) { + // NOTES: dynamic UMMA N writes only `effective_m` columns, so compile-time overlap + // columns beyond the available TMEM loads were never produced + const auto num_loads_before_arrive = cute::min(kNumOverlapLoads, num_stores * kNumTmemLoads); + if (s * kNumTmemLoads + i + 1 == num_loads_before_arrive) { + ptx::tcgen05_before_thread_sync(); + tmem_overlap_barrier->arrive(0u); + } + } + // Notify tensor memory empty (only at the leader CTA) as soon as all loads are in registers + if (s == num_stores - 1 and i == kNumTmemLoads - 1) { + ptx::tcgen05_before_thread_sync(); + tmem_empty_barrier->arrive(0u); + } + + // Store into shared memory + if constexpr (kWithOutputSF) { + // Round the accumulator into BF16 pairs first, so the output bitwise matches + // a BF16 D followed by the standalone cast (power-of-two scaling of BF16 is exact) + // NOTES: with the `.16x256b` load, every pair holds this lane's row pair at one + // of its 4 columns, so all packed multiplications share one per-half sf_inv + nv_bfloat162 values_bf16x2[kNumSwizzleAtomRows / 2]; + #pragma unroll + for (uint32_t pair_idx = 0; pair_idx < kNumSwizzleAtomRows / 2; ++ pair_idx) + values_bf16x2[pair_idx] = math::cast_into_bf16x2(values[pair_idx * 2], values[pair_idx * 2 + 1]); + + // Reduce the row pair's amax and compute the SFs (the SF exponent floor covers + // the amax clamp): the 8 lanes sharing `lane_idx % 4` cover the pair's two SF + // groups, so a packed tree plus a cross-group partial reduction is enough + const auto amax_pair = math::warp_reduce_max<4, true>( + __hmax2(__hmax2(__habs2(values_bf16x2[0]), __habs2(values_bf16x2[1])), + __hmax2(__habs2(values_bf16x2[2]), __habs2(values_bf16x2[3])))); + const auto sf_exp_lower = math::get_ue8m0_sf_exp(amax_pair.x); + const auto sf_exp_upper = math::get_ue8m0_sf_exp(amax_pair.y); + const auto sf_inv = nv_bfloat162(math::get_ue8m0_sf_inv(sf_exp_lower), + math::get_ue8m0_sf_inv(sf_exp_upper)); + + // Cast into FP8 and store with a transposing STSM: the loaded fragment matches + // `stmatrix.m16n8.trans.b8`, with one 8x16 matrix per 16 columns + const auto smem_ptr = smem_base_ptr + (lane_idx % 8) * kSwizzleCDMode + + ((epilogue_warp_idx * 2 + lane_idx / 8) ^ (lane_idx % 8)) * kNumBankGroupBytes; + ptx::SM100_U8x8_STSM_T::copy( + math::scale_bf16x2_into_fp8x4(values_bf16x2[0], values_bf16x2[1], sf_inv, sf_inv), + math::scale_bf16x2_into_fp8x4(values_bf16x2[2], values_bf16x2[3], sf_inv, sf_inv), + smem_ptr); + + // Store the SF byte into the packed word + // NOTES: lanes 0-3 own their row pair's even row, lanes 4-7 the odd row + if (lane_idx < kNumSwizzleAtomRows) + epilogue_op.store_sf(base_m_idx + store_idx * STORE_BLOCK_M + load_idx * kNumSwizzleAtomRows + + 2 * (lane_idx % 4) + lane_idx / 4, + base_n_idx + epilogue_warp_idx * kSFGranN, batch_idx, + static_cast(lane_idx < 4 ? sf_exp_lower : sf_exp_upper)); + } else if constexpr (cute::is_same_v) { uint32_t col = lane_idx / 4; #pragma unroll @@ -83,15 +177,6 @@ sm100_store_cd_swap_ab(const utils::PatternVisitor& smem_cd, uint3 ptx::st_shared(reinterpret_cast(smem_ptr), values[row]); } } else { - // Load from TMEM using `.16x256b` shape to satisfy STSM layout requirements - // Start from lane index 0 - cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr, - values[0], values[1], values[2], values[3]); - // Start from lane index 16 - cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr | 0x00100000, - values[4], values[5], values[6], values[7]); - cutlass::arch::fence_view_async_tmem_load(); - // Destination shared memory address uint32_t row = lane_idx % 8; uint32_t col = (epilogue_warp_idx % 2) * 4 + lane_idx / 8; @@ -107,25 +192,18 @@ sm100_store_cd_swap_ab(const utils::PatternVisitor& smem_cd, uint3 } } - // Notify tensor memory empty (only at the leader CTA) arrival ASAP - // NOTES: only the last stage needs to do this - if (s == num_stores - 1) { - ptx::tcgen05_before_thread_sync(); - tmem_empty_barrier->arrive(0u); - } - // Synchronize all threads and issue TMA cute::tma_store_fence(); cutlass::arch::NamedBarrier::sync(kNumUMMAStoreThreads, 0); if (epilogue_warp_idx == 0 and cute::elect_one_sync()) { #pragma unroll for (uint32_t i = 0; i < STORE_BLOCK_N / STORE_BLOCK_N_ATOM; ++ i) { - auto smem_ptr = smem_cd[tma_stage_idx] + i * STORE_BLOCK_M * STORE_BLOCK_N_ATOM; - uint32_t m_idx = base_m_idx + s * STORE_BLOCK_M; - uint32_t n_idx = epilogue_type_t::apply_index_n(base_n_idx + i * STORE_BLOCK_N_ATOM); + auto smem_ptr = smem.cd[tma_stage_idx] + i * STORE_BLOCK_M * STORE_BLOCK_N_ATOM; + uint32_t m_idx = base_m_idx + store_idx * STORE_BLOCK_M; + uint32_t n_idx = epilogue_op_t::apply_index_n(base_n_idx + i * STORE_BLOCK_N_ATOM); // Issue 2D or 3D TMA store - if constexpr (kGemmType == GemmType::Batched) { + if constexpr (kGemmType == GemmType::Batched or is_k_grouped_contiguous(kGemmType)) { using cute_tma_t = cute::conditional_t; cute_tma_t::copy(&tensor_map_cd, smem_ptr, n_idx, m_idx, batch_idx); diff --git a/deep_gemm/include/deep_gemm/epilogue/transform.cuh b/deep_gemm/include/deep_gemm/epilogue/transform.cuh index 0266f4d402..b9e1bd169f 100644 --- a/deep_gemm/include/deep_gemm/epilogue/transform.cuh +++ b/deep_gemm/include/deep_gemm/epilogue/transform.cuh @@ -1,14 +1,36 @@ #pragma once #include +#include namespace deep_gemm::epilogue::transform { -struct EpilogueIdentity { +// Epilogue operators add behavior (never state) to `EpilogueArgs`, so kernels take +// the host-constructed operator directly as a kernel argument +// NOTES: the operators do not compose with each other +struct EpilogueIdentity: EpilogueArgs { template CUTLASS_DEVICE static uint32_t apply_index_n(const uint32_t& n_idx) { return n_idx; } + + template + CUTLASS_DEVICE void apply_values(uint32_t (&)[kNumValues]) const {} +}; + +// Scale only the product term of a BLAS-style linear combination. Supporting an +// arbitrary beta requires a separate C load/initialization path and does not belong +// in this value-only transform. +struct EpilogueWithAlpha: EpilogueIdentity { + template + CUTLASS_DEVICE void apply_values(uint32_t (&values)[kNumValues]) const { + DG_STATIC_ASSERT(kNumValues % 2 == 0, "Alpha scaling requires float2 alignment"); + const auto values_f32x2 = reinterpret_cast(values); + const auto alpha_f32x2 = make_float2(alpha, alpha); + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumValues / 2; ++ value_idx) + values_f32x2[value_idx] = __fmul2_rn(values_f32x2[value_idx], alpha_f32x2); + } }; template @@ -21,4 +43,33 @@ struct EpilogueHeadSplits: EpilogueIdentity { } }; +// Cast D into FP8 with dynamic per-row, per-32-column UE8M0 SFs, packed into `uint32_t` +// words in a TMA-aligned MN-major layout (the same layout accepted for SFA) +// NOTES: the accumulator is rounded into BF16 before the amax/scale/cast steps, so the +// output bitwise matches a BF16 D followed by the standalone per-token cast kernel +struct EpilogueDynamicScaledFP8: EpilogueIdentity { + static constexpr uint32_t kSFGranN = 32; + + // Store one SF byte (`uint8_t`), or all four SF bytes of one packed word at once + // (`uint32_t`, fully coalesced; only valid when `shape_n % 128 == 0`, so that a word + // never crosses a batch or shape boundary) + // NOTES: batched GEMMs flatten the SF columns of all batches along `(batch_idx, n)`, + // matching a D viewed by the consumer as a `(shape_m, num_batches * shape_n)` 2D tensor + template + CUTLASS_DEVICE void store_sf(const uint32_t& row_idx, const uint32_t& group_n_idx, + const uint32_t& batch_idx, const sf_t& sf) const { + DG_STATIC_ASSERT(cute::is_same_v or cute::is_same_v, + "The SF must be a single byte or a whole packed word"); + if (row_idx >= shape_m or group_n_idx >= shape_n) + return; + const auto sf_idx = (batch_idx * shape_n + group_n_idx) / kSFGranN; + const auto sf_word_ptr = sfd + (sf_idx / 4) * sfd_stride + row_idx; + if constexpr (cute::is_same_v) { + *sf_word_ptr = sf; + } else { + reinterpret_cast(sf_word_ptr)[sf_idx % 4] = sf; + } + } +}; + } // namespace deep_gemm::epilogue::transform diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh index e95965f0e6..20aa4aad99 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh @@ -11,6 +11,7 @@ #include #include #include +#include #include #include #include @@ -29,11 +30,13 @@ template CUTLASS_GLOBAL void __launch_bounds__(kNumNonEpilogueThreads + kNumEpilogueThreads, 1) sm100_bf16_gemm_impl(int* grouped_layout, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + const __grid_constant__ epilogue_op_t epilogue_op, const __grid_constant__ cute::TmaDescriptor tensor_map_a, const __grid_constant__ cute::TmaDescriptor tensor_map_b, const __grid_constant__ cute::TmaDescriptor tensor_map_cd) { @@ -52,6 +55,10 @@ sm100_bf16_gemm_impl(int* grouped_layout, using Barrier = cutlass::arch::ClusterTransactionBarrier; using Allocator = cute::conditional_t; + // The host launches the epilogue operator directly as a kernel argument + DG_STATIC_ASSERT(sizeof(epilogue_op_t) == sizeof(EpilogueArgs), + "Epilogue operators must not add state to `EpilogueArgs`"); + // C/D type: BF16 and FP32 are supported, with or without accumulation DG_STATIC_ASSERT(cute::is_same_v or cute::is_same_v, "Invalid C/D data dtype"); @@ -64,7 +71,7 @@ sm100_bf16_gemm_impl(int* grouped_layout, constexpr uint32_t LOAD_BLOCK_N = BLOCK_N / (kIsMulticastOnA ? 1 : kNumMulticast); DG_STATIC_ASSERT(BLOCK_K_ == 64, "Invalid block K"); DG_STATIC_ASSERT(BLOCK_K % UMMA_K == 0, "Block K must be divisible by UMMA K"); - DG_STATIC_ASSERT(kKAlignment % UMMA_K == 0, "K alignment must be divisible by UMMA K"); + DG_STATIC_ASSERT(not is_k_grouped_contiguous(kGemmType) or kKAlignment % BLOCK_K == 0, "K alignment must be divisible by block K"); DG_STATIC_ASSERT(kNumMulticast == 1 or kNumMulticast == 2, "Only support 1/2 multicast"); DG_STATIC_ASSERT((kSwapAB and BLOCK_N == LAYOUT_AD_M) or (not kSwapAB and (BLOCK_M == 32 or BLOCK_M == 64 or BLOCK_M == LAYOUT_AD_M)), "Invalid block size"); @@ -80,18 +87,8 @@ sm100_bf16_gemm_impl(int* grouped_layout, constexpr uint32_t kNumUMMAStoreThreads = kSwapAB ? kNumEpilogueThreads: STORE_BLOCK_M; DG_STATIC_ASSERT(kNumUMMAStoreThreads % 32 == 0, "Invalid store block M"); - // Share memory sizes - constexpr uint32_t SMEM_CD_SIZE_PER_STAGE = STORE_BLOCK_M * STORE_BLOCK_N * sizeof(cd_dtype_t); - constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_SIZE_PER_STAGE * kNumTMAStoreStages; - constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(cutlass::bfloat16_t); - constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(cutlass::bfloat16_t); - DG_STATIC_ASSERT(SMEM_CD_SIZE % 1024 == 0 and SMEM_A_SIZE_PER_STAGE % 1024 == 0 and SMEM_B_SIZE_PER_STAGE % 1024 == 0, - "Shared memory of A/B must be aligned to 1024 bytes"); - DG_STATIC_ASSERT(kNumTMAStoreStages >= 1, "Invalid number of TMA stages"); - // NOTES: Make sure we have enough shared memory for UMMA padding static constexpr uint32_t UMMA_A_SIZE_PER_STAGE = math::constexpr_align(LOAD_BLOCK_M, LAYOUT_AD_M) * BLOCK_K * sizeof(nv_bfloat16); - DG_STATIC_ASSERT(UMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE * kNumStages, "Memory out of bound for UMMA"); // Real tensor memory size and offsets constexpr uint32_t kNumAccumTmemCols = kNumEpilogueStages * UMMA_N; @@ -118,56 +115,37 @@ sm100_bf16_gemm_impl(int* grouped_layout, shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n; shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k; - // Align to 1024 bytes for swizzle-128B + using SharedStorage = layout::SM100BF16GemmSharedStorage< + kNumStages, kNumEpilogueStages, kNumTMAStoreStages, LOAD_BLOCK_M, LOAD_BLOCK_N, BLOCK_K, + STORE_BLOCK_M, STORE_BLOCK_N, cd_dtype_t>; extern __shared__ __align__(1024) uint8_t smem_buffer[]; - - // D/A/B shared memory - auto smem_cd = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + i * SMEM_CD_SIZE_PER_STAGE); - }); - auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); - }); - auto smem_b = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); - }); - - // Fill barriers - auto barrier_start_ptr = reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE)); - auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (i); }); - auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages + i); }); - auto tmem_full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 2 + i); }); - auto tmem_empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 2 + kNumEpilogueStages + i); }); - auto tensor_core_full_barrier = barrier_start_ptr + kNumStages * 3 + kNumEpilogueStages * 2; - - // Fill the tensor memory pointer - auto tmem_ptr_in_smem = reinterpret_cast(barrier_start_ptr + kNumStages * 3 + kNumEpilogueStages * 2 + 1); - DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); + auto& smem = *reinterpret_cast(smem_buffer); + DG_STATIC_ASSERT(UMMA_A_SIZE_PER_STAGE <= sizeof(smem.a[0]) + sizeof(smem.b), "Memory out of bound for UMMA"); // Initialize barriers if (warp_idx == 1 and cute::elect_one_sync()) { #pragma unroll for (uint32_t i = 0; i < kNumStages; ++ i) { // Arrive only at the leader CTA - full_barriers[i]->init(kNumMulticast); + smem.full_barriers[i].init(kNumMulticast); // Arrive at all CTAs - empty_barriers[i]->init(1); + smem.empty_barriers[i].init(1); } #pragma unroll for (uint32_t i = 0; i < kNumEpilogueStages; ++ i) { // Arrive at all CTAs - tmem_full_barriers[i]->init(1); + smem.tmem_full_barriers[i].init(1); // Arrive only at the leader CTA - tmem_empty_barriers[i]->init(kNumMulticast * kNumUMMAStoreThreads); + smem.tmem_empty_barriers[i].init(kNumMulticast * kNumUMMAStoreThreads); } if constexpr (kTensorCoreUtilControl < 100) - tensor_core_full_barrier->init(1); + smem.tensor_core_full_barrier.init(1); // Make initialized barrier visible in async proxy cutlass::arch::fence_barrier_init(); } else if (warp_idx == 2) { // Allocate tensor memory - Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); + Allocator().allocate(kNumTmemCols, &smem.tmem_ptr); } kNumMulticast > 1 ? comm::cluster_sync_with_relaxed_arrive() : __syncthreads(); @@ -199,10 +177,10 @@ sm100_bf16_gemm_impl(int* grouped_layout, const auto load_block_m = kSwapAB ? scheduler.get_aligned_effective_m_in_block(m_block_idx) / kNumMulticast : LOAD_BLOCK_M; // For k-grouped layout, the number of block K is variable - const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + const auto num_total_k_blocks = cute::max(1u, math::ceil_div(scheduler.current_shape_k, BLOCK_K)); for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { // Wait consumer release - empty_barriers[stage_idx]->wait(phase ^ 1); + smem.empty_barriers[stage_idx].wait(phase ^ 1); // Compute offsets // NOTES: the group is always concatenated with the outer dimension @@ -232,23 +210,23 @@ sm100_bf16_gemm_impl(int* grouped_layout, const uint32_t batch_idx = (kIsBatchedMM ? scheduler.current_group_idx : 0); if constexpr (kMajorA == cute::UMMA::Major::K) tma::copy( - &tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], k_a_idx, m_idx, kNumMulticast, batch_idx); + &tensor_map_a, &smem.full_barriers[stage_idx], smem.a[stage_idx], k_a_idx, m_idx, kNumMulticast, batch_idx); if constexpr (kMajorA == cute::UMMA::Major::MN) tma::copy( - &tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], m_idx, k_a_idx, kNumMulticast, batch_idx); + &tensor_map_a, &smem.full_barriers[stage_idx], smem.a[stage_idx], m_idx, k_a_idx, kNumMulticast, batch_idx); if constexpr (kMajorB == cute::UMMA::Major::K) tma::copy( - &tensor_map_b, full_barriers[stage_idx], smem_b[stage_idx], k_b_idx, n_idx, kNumMulticast, batch_idx); + &tensor_map_b, &smem.full_barriers[stage_idx], smem.b[stage_idx], k_b_idx, n_idx, kNumMulticast, batch_idx); if constexpr (kMajorB == cute::UMMA::Major::MN) tma::copy( - &tensor_map_b, full_barriers[stage_idx], smem_b[stage_idx], n_idx, k_b_idx, kNumMulticast, batch_idx); + &tensor_map_b, &smem.full_barriers[stage_idx], smem.b[stage_idx], n_idx, k_b_idx, kNumMulticast, batch_idx); // Arrive at full barriers - constexpr uint32_t kNumArrivalBytes = SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE; + constexpr uint32_t kNumArrivalBytes = sizeof(smem.a[0]) + sizeof(smem.b[0]); if (is_leader_cta) { - full_barriers[stage_idx]->arrive_and_expect_tx(kNumArrivalBytes * kNumMulticast); + smem.full_barriers[stage_idx].arrive_and_expect_tx(kNumArrivalBytes * kNumMulticast); } else { - full_barriers[stage_idx]->arrive(0u); + smem.full_barriers[stage_idx].arrive(0u); } } } @@ -264,10 +242,10 @@ sm100_bf16_gemm_impl(int* grouped_layout, DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); // Merged stages only happens in NT normal GEMM cases constexpr uint32_t BLOCK_ATOM_K = BLOCK_K / kNumStagesPerMerge; - auto a_desc = mma::sm100::make_umma_desc(smem_a[0], 0, 0); - auto b_desc = mma::sm100::make_umma_desc(smem_b[0], 0, 0); - uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * SMEM_A_SIZE_PER_STAGE / 16 : 0u; - uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 : 0u; + auto a_desc = mma::sm100::make_umma_desc(smem.a[0], 0, 0); + auto b_desc = mma::sm100::make_umma_desc(smem.b[0], 0, 0); + uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * sizeof(smem.a[0]) / 16 : 0u; + uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * sizeof(smem.b[0]) / 16 : 0u; // Checks for MMA instructions // NOTES: CUTLASS does not have such checks except the MMA traits, but we are not using these traits @@ -281,7 +259,7 @@ sm100_bf16_gemm_impl(int* grouped_layout, // Wait tensor memory empty barrier arrival auto accum_stage_idx = scheduler.current_iter % kNumEpilogueStages; auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; - tmem_empty_barriers[accum_stage_idx]->wait(accum_phase_idx ^ 1); + smem.tmem_empty_barriers[accum_stage_idx].wait(accum_phase_idx ^ 1); ptx::tcgen05_after_thread_sync(); // UMMA and empty barrier arrival alias @@ -294,11 +272,11 @@ sm100_bf16_gemm_impl(int* grouped_layout, } }; auto empty_barrier_arrive = [&](const bool& do_tmem_full_arrive) { - umma_arrive(reinterpret_cast(empty_barriers[stage_idx])); + umma_arrive(reinterpret_cast(&smem.empty_barriers[stage_idx])); // NOTES: the tensor memory accumulator pipeline has nothing to do with multicasting if (do_tmem_full_arrive) - umma_arrive(reinterpret_cast(tmem_full_barriers[accum_stage_idx])); + umma_arrive(reinterpret_cast(&smem.tmem_full_barriers[accum_stage_idx])); __syncwarp(); }; @@ -309,11 +287,10 @@ sm100_bf16_gemm_impl(int* grouped_layout, } // Launch MMAs - const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); - constexpr bool kMayHaveTailKBlock = is_k_grouped_contiguous(kGemmType) ? (kKAlignment % BLOCK_K != 0) : (SHAPE_K == 0 or SHAPE_K % BLOCK_K != 0); + const auto num_total_k_blocks = cute::max(1u, math::ceil_div(scheduler.current_shape_k, BLOCK_K)); for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { // Wait TMA arrival - full_barriers[stage_idx]->wait(phase); + smem.full_barriers[stage_idx].wait(phase); ptx::tcgen05_after_thread_sync(); // Issue UMMA in the leader CTA @@ -322,43 +299,21 @@ sm100_bf16_gemm_impl(int* grouped_layout, const auto a_desc_base_lo = __shfl_sync(0xffffffff, a_desc_lo, static_cast(stage_idx)); const auto b_desc_base_lo = __shfl_sync(0xffffffff, b_desc_lo, static_cast(stage_idx)); if (cute::elect_one_sync()) { - auto issue_umma = [&]() { - constexpr uint32_t kAtomKIdx = kUMMAKIdx * UMMA_K / BLOCK_ATOM_K; - constexpr uint32_t kInnerKIdx = kUMMAKIdx * UMMA_K % BLOCK_ATOM_K; + #pragma unroll + for (uint32_t umma_k_idx = 0; umma_k_idx < BLOCK_K / UMMA_K; ++ umma_k_idx) { + const uint32_t atom_k_idx = umma_k_idx * UMMA_K / BLOCK_ATOM_K; + const uint32_t inner_k_idx = umma_k_idx * UMMA_K % BLOCK_ATOM_K; a_desc.lo = mma::sm100::advance_umma_desc_lo( - a_desc_base_lo, kAtomKIdx * LOAD_BLOCK_M * BLOCK_ATOM_K, kInnerKIdx); + a_desc_base_lo, atom_k_idx * LOAD_BLOCK_M * BLOCK_ATOM_K, inner_k_idx); b_desc.lo = mma::sm100::advance_umma_desc_lo( - b_desc_base_lo, kAtomKIdx * LOAD_BLOCK_N * BLOCK_ATOM_K, kInnerKIdx); + b_desc_base_lo, atom_k_idx * LOAD_BLOCK_N * BLOCK_ATOM_K, inner_k_idx); if (kSwapAB) { mma_t::fma(b_desc, a_desc, accum_stage_idx * UMMA_N, - kUMMAKIdx > 0 or k_block_idx > 0, runtime_instr_desc); + umma_k_idx > 0 or k_block_idx > 0, runtime_instr_desc); } else { mma_t::fma(a_desc, b_desc, accum_stage_idx * UMMA_N, - kUMMAKIdx > 0 or k_block_idx > 0, runtime_instr_desc); - } - }; - auto issue_full_k_block = [&]() { - utils::for_each_static_until(std::make_integer_sequence(), issue_umma); - }; - - if constexpr (kMayHaveTailKBlock) { - auto issue_tail_k_block = [&](const uint32_t& remaining_k) { - const auto num_valid_umma_k = math::ceil_div(remaining_k, UMMA_K); - // Prefix expansion uses switch only for small cases to avoid long SASS. - utils::for_each_static_prefix(std::make_integer_sequence(), num_valid_umma_k, issue_umma); - }; - const auto is_last_k_block = k_block_idx == num_total_k_blocks - 1; - if (is_last_k_block) { - const auto remaining_k = scheduler.current_shape_k - k_block_idx * BLOCK_K; - if (remaining_k < BLOCK_K) - issue_tail_k_block(remaining_k); - else - issue_full_k_block(); - } else { - issue_full_k_block(); + umma_k_idx > 0 or k_block_idx > 0, runtime_instr_desc); } - } else { - issue_full_k_block(); } } __syncwarp(); @@ -371,11 +326,11 @@ sm100_bf16_gemm_impl(int* grouped_layout, DG_STATIC_ASSERT(kTensorCoreUtilControl > 0, "Invalid tensor utilization control"); if constexpr (kTensorCoreUtilControl < 100) { // For utilization control - umma_arrive(reinterpret_cast(tensor_core_full_barrier)); + umma_arrive(reinterpret_cast(&smem.tensor_core_full_barrier)); __syncwarp(); // Wait for last UMMA to be done - tensor_core_full_barrier->wait(tensor_core_phase); + smem.tensor_core_full_barrier.wait(tensor_core_phase); tensor_core_phase ^= 1; // Sleep for certain cycles @@ -393,7 +348,7 @@ sm100_bf16_gemm_impl(int* grouped_layout, const auto iter_idx = scheduler.current_iter - 1; if (kNumMulticast > 1 and iter_idx >= 0) { const auto accum_phase_idx = (iter_idx / kNumEpilogueStages) & 1; - tmem_empty_barriers[iter_idx % kNumEpilogueStages]->wait(accum_phase_idx); + smem.tmem_empty_barriers[iter_idx % kNumEpilogueStages].wait(accum_phase_idx); } } else if (warp_idx >= kNumNonEpilogueThreads / 32 and warp_idx < (kNumNonEpilogueThreads + kNumUMMAStoreThreads) / 32) { // Epilogue warp groups @@ -402,7 +357,7 @@ sm100_bf16_gemm_impl(int* grouped_layout, // NOTES: tensor memory addresses are simplified, as the hardware will ignore the warp index bits, // i.e., no need for `tmem_ptr |= (epilogue_warp_idx * 32) << 16`. // NOTES: we also forbid two CTAs to share the same SM and its tensor memory - DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(tmem_ptr_in_smem) == 0); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr) == 0); // Share store pipeline between blocks uint32_t tma_stage_idx = 0; @@ -413,36 +368,47 @@ sm100_bf16_gemm_impl(int* grouped_layout, auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; // Wait UMMA arrival - tmem_full_barriers[accum_stage_idx]->wait(accum_phase_idx); + smem.tmem_full_barriers[accum_stage_idx].wait(accum_phase_idx); ptx::tcgen05_after_thread_sync(); // Load from tensor memory into registers, and write shared memory with STSM const auto tmem_base_addr = accum_stage_idx * UMMA_N; + // Whether the group offset is encoded in the flattened CD M coordinate + constexpr bool kCDWithGroupOffset = not is_m_grouped_contiguous(kGemmType) and not is_k_grouped_contiguous(kGemmType); const auto base_m_idx = scheduler.template get_global_idx< - (not is_m_grouped_contiguous(kGemmType)), sched::IndexType::MN>(shape_m, BLOCK_M, m_block_idx); + kCDWithGroupOffset, sched::IndexType::MN>(shape_m, BLOCK_M, m_block_idx); const auto base_n_idx = n_block_idx * BLOCK_N; + const bool is_empty_group = is_k_grouped_contiguous(kGemmType) and scheduler.current_shape_k == 0; if constexpr (kSwapAB) { const auto effective_m = scheduler.get_aligned_effective_m_in_block(m_block_idx); epilogue::sm100_store_cd_swap_ab - (smem_cd, tma_stage_idx, tmem_base_addr, + 0, + kGemmType, kWithAccumulation> + (smem, tma_stage_idx, tmem_base_addr, base_m_idx, base_n_idx, scheduler.current_group_idx, + is_empty_group, effective_m, epilogue_warp_idx, lane_idx, - tmem_empty_barriers[accum_stage_idx], + epilogue_op, + false, + &smem.tmem_empty_barriers[accum_stage_idx], + &smem.tmem_empty_barriers[accum_stage_idx], tensor_map_cd); } else { epilogue::sm100_store_cd - (smem_cd, tma_stage_idx, tmem_base_addr, + 0, + kGemmType, kWithAccumulation> + (smem, tma_stage_idx, tmem_base_addr, base_m_idx, base_n_idx, scheduler.current_group_idx, + is_empty_group, epilogue_warp_idx, lane_idx, - tmem_empty_barriers[accum_stage_idx], + epilogue_op, + false, + &smem.tmem_empty_barriers[accum_stage_idx], + &smem.tmem_empty_barriers[accum_stage_idx], tensor_map_cd); } } diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_gate.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_gate.cuh new file mode 100644 index 0000000000..e6d3636033 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_gate.cuh @@ -0,0 +1,483 @@ +#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 + +namespace deep_gemm { + +namespace mega_gate_layout = layout::mega_gate; + +CUTLASS_HOST_DEVICE constexpr uint32_t get_num_stages_per_merge(const uint32_t num_stages, + const uint32_t num_k_atoms_per_slice) { + return num_stages / 3 >= 8 and num_k_atoms_per_slice % 3 == 0 ? 3u : + num_stages / 2 >= 8 and num_k_atoms_per_slice % 2 == 0 ? 2u : 1u; +} + +template +CUTLASS_GLOBAL void __launch_bounds__(mega_gate_layout::kNumNonEpilogueThreads + kNumGateThreads, 1) +sm100_bf16_mega_gate_impl(const __grid_constant__ cute::TmaDescriptor tensor_map_x, + const __grid_constant__ cute::TmaDescriptor tensor_map_weight, + const float* bias, const float* image_bias, const bool* image_token_mask, + const bool* mask, const int* to_physical_map, const int* logical_count, + int64_t* topk_idx, int64_t* unmapped_topk_idx, float* topk_weights, + void* gmem_scratch, void* gmem_score_barriers, + uint32_t num_tokens, uint32_t num_routed_experts, + uint32_t num_shared_experts, uint32_t num_duplicate_experts, + float routed_scaling_factor, int ep_rank, int64_t unmapped_topk_idx_stride, + const bool* fix_routing_mask, const bool* force_random) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using Allocator = cute::conditional_t; + + constexpr uint32_t kNumNonEpilogueWarps = mega_gate_layout::kNumNonEpilogueThreads / 32; + constexpr uint32_t kNumGateWarps = kNumGateThreads / 32; + constexpr uint32_t BLOCK_ATOM_K = mega_gate_layout::BLOCK_K; + constexpr uint32_t kNumExpertsPerGroup = kNumRoutedExperts / kNumExpertGroups; + + constexpr uint32_t UMMA_M = kNumExpertsPerGroup; + constexpr uint32_t UMMA_N = BLOCK_TOKENS; + constexpr uint32_t kNumLogicalCtas = kNumMmaCtas * kNumExpertGroups * kNumSplitK; + + constexpr uint32_t kNumStagesPerMerge = get_num_stages_per_merge(kNumStages_, (SHAPE_K / BLOCK_ATOM_K) / kNumSplitK); + constexpr uint32_t BLOCK_K = BLOCK_ATOM_K * kNumStagesPerMerge; + constexpr uint32_t kNumStages = kNumStages_ / kNumStagesPerMerge; + constexpr uint32_t LOAD_BLOCK_M = UMMA_N / kNumMmaCtas; + constexpr uint32_t LOAD_BLOCK_N = UMMA_M / kNumMmaCtas; + DG_STATIC_ASSERT(kNumExpertsPerGroup <= 256, "An expert group must be a single UMMA tile"); + constexpr uint32_t kNumEpilogueStages = mega_gate_layout::kNumEpilogueStages; + DG_STATIC_ASSERT(UMMA_N <= mega_gate_layout::kNumMaxBlockTokens, "Block tokens exceed a single UMMA tile"); + using SharedStorage = mega_gate_layout::SharedStorage; + using Barrier = typename SharedStorage::Barrier; + + constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); + + const auto mma_cta_rank = cute::block_rank_in_cluster(); + const auto mma_cluster_idx = blockIdx.x / kNumMmaCtas; + const auto slice_idx = mma_cluster_idx % (kNumExpertGroups * kNumSplitK); + const auto expert_group_idx = slice_idx % kNumExpertGroups; + const auto split_k_idx = slice_idx / kNumExpertGroups; + const auto logical_cta_rank = slice_idx * kNumMmaCtas + mma_cta_rank; + constexpr uint32_t kNumKBlocksPerSlice = (SHAPE_K / BLOCK_K) / kNumSplitK; + const auto split_k_offset = split_k_idx * kNumKBlocksPerSlice * BLOCK_K; + const auto is_leader_cta = mma_cta_rank == 0; + const auto warp_idx = cutlass::canonical_warp_idx_sync(); + const auto lane_idx = ptx::get_lane_idx(); + + if (warp_idx == 0) { + cute::prefetch_tma_descriptor(&tensor_map_x); + cute::prefetch_tma_descriptor(&tensor_map_weight); + } + + extern __shared__ __align__(mega_gate_layout::kSharedMemoryAlignment) uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + const auto smem_x = smem.x; + const auto smem_weight = smem.weight; + const auto smem_bias = smem.get_bias_ptr(); + const auto smem_image_bias = smem.get_image_bias_ptr(); + const auto smem_logical_count = smem.get_logical_count_ptr(); + const auto routed_logical_count = SharedStorage::kCacheLogicalCount ? smem_logical_count : logical_count; + + if constexpr (kNumMmaCtas > 1) + comm::cluster_sync_with_relaxed_arrive(); + + if (warp_idx < 4 and cute::elect_one_sync()) { + #pragma unroll + for (uint32_t stage_idx = warp_idx; stage_idx < kNumStages; stage_idx += 4) { + smem.full_barriers[stage_idx].init(kNumMmaCtas); + smem.empty_barriers[stage_idx].init(1); + } + if (warp_idx == 0) { + #pragma unroll + for (uint32_t accum_stage_idx = 0; accum_stage_idx < kNumEpilogueStages; ++ accum_stage_idx) { + smem.tmem_full_barriers[accum_stage_idx].init(1); + smem.tmem_empty_barriers[accum_stage_idx].init(kNumMmaCtas); + } + if constexpr (SharedStorage::kHasMetadataCache) + smem.metadata_ready_barrier.init(1); + } + cutlass::arch::fence_barrier_init(); + } + __syncwarp(); + if (warp_idx == 2) + Allocator().allocate(kNumTmemCols, &smem.tmem_ptr); + comm::cluster_sync_with_relaxed_arrive(); + + cudaGridDependencySynchronize(); + + const auto get_tmem_base_addr = [](const uint32_t& accum_stage_idx) { + return accum_stage_idx * UMMA_N; + }; + const sched::mega_gate::Scheduler scheduler(num_tokens); + const mega_gate_layout::Workspace workspace(gmem_scratch, + gmem_score_barriers); + const auto get_effective_umma_n = [&](const uint32_t& token_block_idx) { + return math::align(cute::min(BLOCK_TOKENS, num_tokens - token_block_idx * BLOCK_TOKENS), 16u); + }; + + const auto produce_tma = [&](const uint32_t& token_block_idx, const uint32_t& iter_idx) { + const auto first_k_block_idx = iter_idx * kNumKBlocksPerSlice; + const auto expert_base_idx = expert_group_idx * kNumExpertsPerGroup + mma_cta_rank * LOAD_BLOCK_N; + const auto token_base_idx = token_block_idx * BLOCK_TOKENS; + const auto effective_umma_n = get_effective_umma_n(token_block_idx); + const auto x_token_idx = token_base_idx + mma_cta_rank * (effective_umma_n / kNumMmaCtas); + if (cute::elect_one_sync()) { + #pragma unroll 4 + for (uint32_t k_block_idx = 0; k_block_idx < kNumKBlocksPerSlice; ++ k_block_idx) { + const auto pipeline_idx = first_k_block_idx + k_block_idx; + const auto stage_idx = pipeline_idx % kNumStages; + const auto phase = pipeline_idx / kNumStages & 1; + smem.empty_barriers[stage_idx].wait(phase ^ 1); + + const auto k_idx = split_k_offset + k_block_idx * BLOCK_K; + tma::copy(cute::TMA::CacheHintSm100::EVICT_FIRST)>(&tensor_map_x, + &smem.full_barriers[stage_idx], + smem_x[stage_idx], k_idx, + x_token_idx, kNumMmaCtas); + tma::copy(&tensor_map_weight, + &smem.full_barriers[stage_idx], + smem_weight[stage_idx], k_idx, + expert_base_idx, kNumMmaCtas); + if (is_leader_cta) { + constexpr uint32_t kNumArrivalBytes = sizeof(smem_x[0]) + sizeof(smem_weight[0]); + smem.full_barriers[stage_idx].arrive_and_expect_tx(kNumArrivalBytes * kNumMmaCtas); + } else { + smem.full_barriers[stage_idx].arrive(0u); + } + } + } + }; + + const auto produce_mma = [&](const uint32_t& token_block_idx, const uint32_t& iter_idx) { + const auto first_k_block_idx = iter_idx * kNumKBlocksPerSlice; + auto instr_desc = cute::UMMA::make_instr_desc(); + auto x_desc = mma::sm100::make_umma_desc(smem_x[0], 0, 0); + auto weight_desc = mma::sm100::make_umma_desc(smem_weight[0], 0, 0); + + const auto x_desc_lo = lane_idx < kNumStages ? x_desc.lo + lane_idx * sizeof(smem_x[0]) / 16 : 0u; + const auto weight_desc_lo = lane_idx < kNumStages ? weight_desc.lo + lane_idx * sizeof(smem_weight[0]) / 16 : 0u; + + auto umma_arrive = [&](const Barrier* barrier) { + if constexpr (kNumMmaCtas == 1) { + cutlass::arch::umma_arrive(reinterpret_cast(barrier)); + } else { + constexpr auto pair_cta_mask = static_cast((1u << kNumMmaCtas) - 1u); + cutlass::arch::umma_arrive_multicast_2x1SM(reinterpret_cast(barrier), pair_cta_mask); + } + }; + + const auto effective_umma_n = get_effective_umma_n(token_block_idx); + mma::sm100::update_instr_desc_with_umma_n(instr_desc, effective_umma_n); + + const auto accum_stage_idx = iter_idx % kNumEpilogueStages; + const auto accum_phase_idx = iter_idx / kNumEpilogueStages & 1; + smem.tmem_empty_barriers[accum_stage_idx].wait(accum_phase_idx ^ 1); + ptx::tcgen05_after_thread_sync(); + + for (uint32_t k_block_idx = 0; k_block_idx < kNumKBlocksPerSlice; ++ k_block_idx) { + const auto pipeline_idx = first_k_block_idx + k_block_idx; + const auto stage_idx = pipeline_idx % kNumStages; + const auto phase = pipeline_idx / kNumStages & 1; + const auto x_desc_base_lo = ptx::exchange(x_desc_lo, stage_idx); + const auto weight_desc_base_lo = ptx::exchange(weight_desc_lo, stage_idx); + smem.full_barriers[stage_idx].wait(phase); + ptx::tcgen05_after_thread_sync(); + const auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(instr_desc); + + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t umma_k_idx = 0; umma_k_idx < BLOCK_K / 16; ++ umma_k_idx) { + const auto atom_k_idx = umma_k_idx * 16 / BLOCK_ATOM_K; + const auto inner_k_idx = umma_k_idx * 16 % BLOCK_ATOM_K; + using mma_t = cute::conditional_t; + x_desc.lo = mma::sm100::advance_umma_desc_lo( + x_desc_base_lo, atom_k_idx * LOAD_BLOCK_M * BLOCK_ATOM_K, inner_k_idx); + weight_desc.lo = mma::sm100::advance_umma_desc_lo( + weight_desc_base_lo, atom_k_idx * LOAD_BLOCK_N * BLOCK_ATOM_K, + inner_k_idx); + mma_t::fma(weight_desc, x_desc, get_tmem_base_addr(accum_stage_idx), + umma_k_idx > 0 or k_block_idx > 0, runtime_instr_desc); + } + } + __syncwarp(); + + if (k_block_idx == kNumKBlocksPerSlice - 1) + umma_arrive(&smem.tmem_full_barriers[accum_stage_idx]); + __syncwarp(); + + umma_arrive(&smem.empty_barriers[stage_idx]); + } + }; + + if constexpr (SharedStorage::kHasMetadataCache) { + if (warp_idx == kNumNonEpilogueWarps - 1) { + #pragma unroll + for (uint32_t cache_round_idx = 0; cache_round_idx < kNumRoutedExperts / 32; ++ cache_round_idx) { + const auto cache_idx = cache_round_idx * 32 + lane_idx; + const auto is_valid_expert = kFullExpertTile or cache_idx < num_routed_experts; + if constexpr (kHasBias) + smem_bias[cache_idx] = is_valid_expert ? bias[cache_idx] : 0.0f; + if constexpr (kHasImageTokenMask) + smem_image_bias[cache_idx] = is_valid_expert ? image_bias[cache_idx] : 0.0f; + if constexpr (SharedStorage::kCacheLogicalCount) + smem_logical_count[cache_idx] = is_valid_expert ? logical_count[cache_idx] : 0; + } + cutlass::arch::fence_view_async_shared(); + __syncwarp(); + if (lane_idx == 0) + smem.metadata_ready_barrier.arrive(); + } + } + + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr) == 0); + + const auto gate_warp_idx = warp_idx - kNumNonEpilogueWarps; + const auto run_gate = [&](const uint32_t& token_block_idx, const uint32_t& current_iter_idx) { + const auto token_base_idx = token_block_idx * BLOCK_TOKENS; + const auto num_valid_tokens = cute::min(BLOCK_TOKENS, num_tokens - token_block_idx * BLOCK_TOKENS); + const auto effective_umma_n = get_effective_umma_n(token_block_idx); + const auto accum_stage_idx = current_iter_idx % kNumEpilogueStages; + const auto accum_phase_idx = (current_iter_idx / kNumEpilogueStages) & 1; + const auto score_barrier_ptr = workspace.get_score_barrier_ptr(token_block_idx); + + if (gate_warp_idx == kNumGateWarps - 1 and lane_idx == 31 and kNumLogicalCtas > 1) { + if (logical_cta_rank == 0) + sched::mega_gate::ScoreBarrier::init(score_barrier_ptr); + else + sched::mega_gate::ScoreBarrier::wait_init(score_barrier_ptr); + } + + __syncwarp(); + smem.tmem_full_barriers[accum_stage_idx].wait(accum_phase_idx); + ptx::tcgen05_after_thread_sync(); + + const auto global_score_tile = workspace.get_score_ptr(token_block_idx); + const auto split_score_tile = workspace.get_score_ptr(token_block_idx, split_k_idx); + const auto expert_group_base_idx = expert_group_idx * kNumExpertsPerGroup; + constexpr bool kHasScoringBias = kHasBias or kHasImageTokenMask; + constexpr bool kScoreOnStore = kHasScoringBias and kNumSplitK == 1; + auto gate_is_valid = epilogue::mega_gate::sm100_store_tmem_scores_to_global< + kScoreOnStore ? kScoringType : static_cast(mega_gate_layout::ScoringType::Identity), + kNumRoutedExperts, UMMA_M, kNumMmaCtas, kNumGateWarps>( + split_score_tile, + get_tmem_base_addr(accum_stage_idx), + effective_umma_n, expert_group_base_idx, + mma_cta_rank, gate_warp_idx + ); + cutlass::arch::NamedBarrier::sync(kNumGateWarps * 32, 0); + if (gate_warp_idx == kNumGateWarps - 1 and lane_idx == 31) { + ptx::tcgen05_before_thread_sync(); + smem.tmem_empty_barriers[accum_stage_idx].arrive(0u); + if constexpr (kNumLogicalCtas > 1) + sched::mega_gate::ScoreBarrier::arrive(score_barrier_ptr); + } + if (gate_warp_idx == 0 and lane_idx == 0 and kNumLogicalCtas > 1) + sched::mega_gate::ScoreBarrier::wait(score_barrier_ptr); + cutlass::arch::NamedBarrier::sync(kNumGateWarps * 32, 0); + const auto load_scores = [&](const uint32_t& token_idx_in_block, const uint32_t& global_expert_idx, auto& scores) { + constexpr uint32_t kNumValues = sizeof(scores) / sizeof(float); + using vec_t = cute::conditional_t; + const auto score_ptr = global_score_tile + + static_cast(token_idx_in_block) * kNumRoutedExperts + global_expert_idx; + const auto vec_score_ptr = reinterpret_cast(score_ptr); + scores = ptx::ld_global(vec_score_ptr); + #pragma unroll + for (uint32_t split_idx = 1; split_idx < kNumSplitK; ++ split_idx) { + const auto partial_score_ptr = score_ptr + split_idx * UMMA_N * kNumRoutedExperts; + const auto partial_scores = ptx::ld_global(reinterpret_cast(partial_score_ptr)); + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumValues; ++ value_idx) + reinterpret_cast(&scores)[value_idx] += reinterpret_cast(&partial_scores)[value_idx]; + } + }; + const auto load_reduced_score = [&](const uint32_t& token_idx_in_block, const uint32_t& global_expert_idx) { + auto score = 0.0f; + load_scores(token_idx_in_block, global_expert_idx, score); + return score; + }; + for (uint32_t token_idx_in_block = gate_warp_idx * kNumLogicalCtas + logical_cta_rank; + token_idx_in_block < num_valid_tokens; + token_idx_in_block += kNumGateWarps * kNumLogicalCtas) { + const auto global_token_idx = token_base_idx + token_idx_in_block; + if (kMaskExists and not mask[global_token_idx]) { + const auto num_physical_topk = kNumTopk + num_shared_experts; + if (lane_idx < num_physical_topk) { + const auto output_idx = global_token_idx * num_physical_topk + lane_idx; + topk_idx[output_idx] = -1; + topk_weights[output_idx] = 0.0f; + } + if constexpr (kUnmappedTopkIdxExists) { + if (lane_idx < kNumTopk) + unmapped_topk_idx[global_token_idx * unmapped_topk_idx_stride + lane_idx] = -1; + } + continue; + } + if constexpr (kForceRandomExists) { + if (force_random[global_token_idx]) { + epilogue::mega_gate::store_random_topk_token( + global_token_idx, lane_idx, + num_routed_experts, num_shared_experts, + logical_count, ep_rank, + unmapped_topk_idx_stride, topk_idx, + unmapped_topk_idx, topk_weights + ); + continue; + } + } + if constexpr (kFixRoutingMaskExists) { + if (fix_routing_mask[global_token_idx]) { + int selected_expert_idx = -1; + auto selected_unbiased_score = 0.0f; + if (lane_idx < kNumTopk) { + const auto fixed_expert_idx = unmapped_topk_idx[global_token_idx * unmapped_topk_idx_stride + lane_idx]; + DG_TRAP_ONLY_DEVICE_ASSERT(fixed_expert_idx >= 0 and fixed_expert_idx < num_routed_experts); + selected_expert_idx = static_cast(fixed_expert_idx); + selected_unbiased_score = load_reduced_score(token_idx_in_block, static_cast(selected_expert_idx)); + if constexpr (kNumSplitK > 1 or not kHasScoringBias) + selected_unbiased_score = epilogue::mega_gate::apply_scoring(selected_unbiased_score); + } + epilogue::mega_gate::store_topk_token< + kNumTopk, kUnmappedTopkIdxExists, kToPhysicalMapExists>(global_token_idx, lane_idx, + num_routed_experts, + num_shared_experts, + num_duplicate_experts, + routed_scaling_factor, ep_rank, + unmapped_topk_idx_stride, + to_physical_map, logical_count, + routed_logical_count, topk_idx, + unmapped_topk_idx, topk_weights, + selected_expert_idx, + selected_unbiased_score); + continue; + } + } + constexpr uint32_t kNumVectorElements = 4; + constexpr uint32_t kNumExpertsPerWave = 32 * kNumVectorElements; + constexpr uint32_t kNumExpertWaves = kNumRoutedExperts / kNumExpertsPerWave; + constexpr uint32_t kNumExpertsPerLane = kNumRoutedExperts / 32; + float unbiased_scores_local[kNumExpertsPerLane]; + float scores_local[kNumExpertsPerLane]; + const auto is_image_token = kHasImageTokenMask and image_token_mask[global_token_idx]; + #pragma unroll + for (uint32_t expert_wave_idx = 0; expert_wave_idx < kNumExpertWaves; ++ expert_wave_idx) { + const auto global_expert_base_idx = expert_wave_idx * kNumExpertsPerWave + + lane_idx * kNumVectorElements; + const auto local_expert_base_idx = expert_wave_idx * kNumVectorElements; + auto& score_values = *reinterpret_cast(scores_local + local_expert_base_idx); + load_scores(token_idx_in_block, global_expert_base_idx, score_values); + } + #pragma unroll + for (uint32_t expert_wave_idx = 0; expert_wave_idx < kNumExpertWaves; ++ expert_wave_idx) { + const auto global_expert_base_idx = expert_wave_idx * kNumExpertsPerWave + lane_idx * kNumVectorElements; + const auto local_expert_base_idx = expert_wave_idx * kNumVectorElements; + auto& score_values = *reinterpret_cast(scores_local + local_expert_base_idx); + if constexpr (kHasScoringBias and kNumSplitK > 1) + epilogue::mega_gate::apply_scoring(score_values); + const float* bias_ptr = nullptr; + if constexpr (kHasBias) + bias_ptr = smem_bias; + if constexpr (kHasImageTokenMask) + bias_ptr = is_image_token ? smem_image_bias : bias_ptr; + const auto bias_values = bias_ptr + ? ptx::ld_shared(reinterpret_cast(bias_ptr + global_expert_base_idx)) + : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + const auto bias_local = reinterpret_cast(&bias_values); + const auto is_valid_expert = kFullExpertTile or + expert_wave_idx + 1 < kNumExpertWaves or + global_expert_base_idx < num_routed_experts; + #pragma unroll + for (uint32_t value_idx = 0; value_idx < kNumVectorElements; ++ value_idx) { + const auto unbiased_score = score_values[value_idx]; + if constexpr (kNumSplitK > 1) + unbiased_scores_local[local_expert_base_idx + value_idx] = unbiased_score; + scores_local[local_expert_base_idx + value_idx] = is_valid_expert + ? unbiased_score + bias_local[value_idx] + : -cute::numeric_limits::infinity(); + } + } + + int selected_expert_idx = -1; + auto selected_unbiased_score = 0.0f; + epilogue::mega_gate::select_warp_topk(scores_local, 0, selected_expert_idx); + gate_is_valid = gate_is_valid and (lane_idx >= kNumTopk or + (selected_expert_idx >= 0 and selected_expert_idx < static_cast(num_routed_experts))); + selected_expert_idx = lane_idx < kNumTopk + ? cute::min(cute::max(selected_expert_idx, 0), static_cast(num_routed_experts) - 1) + : selected_expert_idx; + if constexpr (kNumSplitK == 1) { + if (lane_idx < kNumTopk) + selected_unbiased_score = load_reduced_score(token_idx_in_block, static_cast(selected_expert_idx)); + } else { + selected_unbiased_score = epilogue::mega_gate::select_warp_expert_value(unbiased_scores_local, 0, selected_expert_idx); + } + if constexpr (not kHasScoringBias) { + if (lane_idx < kNumTopk) + selected_unbiased_score = epilogue::mega_gate::apply_scoring(selected_unbiased_score); + } + + epilogue::mega_gate::store_topk_token< + kNumTopk, kUnmappedTopkIdxExists, kToPhysicalMapExists>(global_token_idx, lane_idx, + num_routed_experts, num_shared_experts, + num_duplicate_experts, routed_scaling_factor, + ep_rank, unmapped_topk_idx_stride, + to_physical_map, logical_count, + routed_logical_count, topk_idx, + unmapped_topk_idx, topk_weights, + selected_expert_idx, + selected_unbiased_score); + } + + DG_TRAP_ONLY_DEVICE_ASSERT(gate_is_valid and "Mega gate scores must be finite and top-k indices valid"); + }; + + if (warp_idx == 0) + scheduler.run(produce_tma); + else if (warp_idx == 1 and is_leader_cta) + scheduler.run(produce_mma); + + if (warp_idx >= kNumNonEpilogueWarps) { + if constexpr (SharedStorage::kHasMetadataCache) + smem.metadata_ready_barrier.wait(0); + scheduler.run(run_gate); + } + + comm::cluster_sync_with_relaxed_arrive(); + __syncwarp(); + if (warp_idx == 0) + Allocator().free(0, kNumTmemCols); +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports SM100-family GPUs"); +#endif +} + +} // namespace deep_gemm + +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh index 9e63ca26a0..bd858d3b5e 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh @@ -107,6 +107,8 @@ sm100_bf16_mega_moe_impl(void* y, ); const auto workspace = buffer.workspace; + using L2KBlockDependency = sched::L2KBlockDependency; + // Data types using a_dtype_t = cutlass::bfloat16_t; using b_dtype_t = cutlass::bfloat16_t; @@ -230,9 +232,10 @@ sm100_bf16_mega_moe_impl(void* y, // Allocate tensor memory Allocator().allocate(kNumTmemCols, &shared_storage.tmem_ptr_in_smem); } - // NOTES: Using `.relaxed` is allowed here since `fence_barrier_init` is `.release.cluster`, - // and `barrier.cluster.wait.aligned` is by default `.acquire` - comm::cluster_sync_with_relaxed_arrive(); + cute::cluster_sync(); + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); // Task scheduler auto scheduler = sched::MegaMoEScheduler< @@ -267,8 +270,7 @@ sm100_bf16_mega_moe_impl(void* y, // NVLink barrier tags constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; - constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; - constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 2; // Adjust registers // NOTES: more experts per rank will cost more schedulers' registers @@ -304,7 +306,7 @@ sm100_bf16_mega_moe_impl(void* y, int expert_idx = -1; if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { expert_idx = static_cast( - __ldg(buffer.input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + buffer.input_topk_idx_buffer.get_base_ptr()[i * kNumTopk + lane_idx]); if (expert_idx >= 0) process(i * kNumTopk + lane_idx, expert_idx); } @@ -344,6 +346,12 @@ sm100_bf16_mega_moe_impl(void* y, // Write expert count if (sm_idx == 0) { + // Push this launch's grid index to every peer for tagging combine readiness; +1 differs from a zeroed workspace + DG_STATIC_ASSERT(kNumRanks <= kNumDispatchThreads, "Insufficient threads for the grid index push"); + if (thread_idx < kNumRanks) + *sym_buffer.map(workspace.get_peer_grid_idx_ptr(sym_buffer.rank_idx), thread_idx) = ptx::get_grid_idx() + 1; + __syncwarp(); + #pragma unroll for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { const auto dst_rank_idx = i / kNumExpertsPerRank; @@ -503,6 +511,10 @@ sm100_bf16_mega_moe_impl(void* y, cute::tma_store_arrive(); ptx::tma_store_wait<0>(); }; + // Load the weight first, so that its remote latency overlaps with the token pull below + const auto weight = *sym_buffer.map( + buffer.input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); if (cute::elect_one_sync()) { #pragma unroll for (uint32_t i = 0; i < kNumChunks; ++ i) { @@ -519,10 +531,6 @@ sm100_bf16_mega_moe_impl(void* y, // Store weights and metadata if (cute::elect_one_sync()) { - // Load weights - const auto weight = *sym_buffer.map( - buffer.input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, - current_rank_in_expert_idx); *buffer.l1_topk_weights_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).template get_base_ptr() = weight; // Write source metadata for combine write-back (logical pool token) @@ -593,7 +601,7 @@ sm100_bf16_mega_moe_impl(void* y, for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { *workspace.get_l1_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; *workspace.get_l1_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; - *workspace.get_l2_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l2_full_mask_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; *workspace.get_l2_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; } __syncwarp(); @@ -631,18 +639,19 @@ sm100_bf16_mega_moe_impl(void* y, const auto ptr = workspace.get_l1_full_count_ptr(block_idx); const auto num_expected_tokens = BLOCK_M * (pool_block_idx / kNumRingBlocks + 1); while (ptx::ld_acq(ptr) != num_expected_tokens); - } else if (task_info.block_phase == sched::BlockPhase::Linear2) { - const auto ptr = workspace.get_l2_full_count_ptr(block_idx); - const auto num_expected_blocks = L2_SHAPE_K / (BLOCK_N / 2) * (pool_block_idx / kNumRingBlocks + 1); - while (ptx::ld_acq(ptr) != num_expected_blocks); } else if (task_info.block_phase == sched::BlockPhase::SharedLinear2) { const auto ptr = workspace.get_shared_l2_full_count_ptr(block_idx); const auto num_expected_blocks = SHARED_L2_SHAPE_K / (BLOCK_N / 2); while (ptx::ld_acq(ptr) != num_expected_blocks); } + L2KBlockDependency l2_k_block_dependency(workspace.get_l2_full_mask_ptr(ring_block_idx), pool_block_idx / kNumRingBlocks); + // Issue TMA for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + if (task_info.block_phase == sched::BlockPhase::Linear2) + l2_k_block_dependency.wait(k_block_idx); + // Wait consumer release shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); @@ -716,11 +725,10 @@ sm100_bf16_mega_moe_impl(void* y, auto instr_desc = cute::UMMA::make_instr_desc(); DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); - uint32_t lane_stage = lane_idx < kNumStages ? lane_idx : 0u; - auto a_desc = mma::sm100::make_umma_desc(shared_storage.smem_a[lane_stage], 0, 0); - auto b_desc = mma::sm100::make_umma_desc(shared_storage.smem_b[lane_stage], 0, 0); - uint32_t a_desc_lo = a_desc.lo; - uint32_t b_desc_lo = b_desc.lo; + auto a_desc = mma::sm100::make_umma_desc(shared_storage.smem_a[0], 0, 0); + auto b_desc = mma::sm100::make_umma_desc(shared_storage.smem_b[0], 0, 0); + const uint32_t a_desc_lo = a_desc.lo; + const uint32_t b_desc_lo = b_desc.lo; // Checks for MMA instructions DG_STATIC_ASSERT((UMMA_M == 64 and UMMA_N % 8 == 0 and 8 <= UMMA_N and UMMA_N <= 256) or @@ -763,8 +771,8 @@ sm100_bf16_mega_moe_impl(void* y, shared_storage.full_barriers[stage_idx].wait(phase); ptx::tcgen05_after_thread_sync(); - const auto a_desc_base_lo = ptx::exchange(a_desc_lo, stage_idx); - const auto b_desc_base_lo = ptx::exchange(b_desc_lo, stage_idx); + const uint32_t a_desc_base_lo = a_desc_lo + stage_idx * sizeof(SharedStorage::smem_a[0]) / 16; + const uint32_t b_desc_base_lo = b_desc_lo + stage_idx * sizeof(SharedStorage::smem_b[0]) / 16; if (cute::elect_one_sync()) { #pragma unroll for (uint32_t umma_k_block_idx = 0; umma_k_block_idx < BLOCK_K / UMMA_BLOCK_K; ++ umma_k_block_idx) { @@ -1001,8 +1009,7 @@ sm100_bf16_mega_moe_impl(void* y, ptx::red_add_rel( workspace.get_shared_l2_full_count_ptr(pool_block_idx), 1u); } else { - ptx::red_add_rel( - workspace.get_l2_full_count_ptr(ring_block_idx), 1u); + L2KBlockDependency::arrive(workspace.get_l2_full_mask_ptr(ring_block_idx), n_block_idx); // Increment L1 empty count for this physical slot (one per N block) ptx::red_add( @@ -1034,6 +1041,18 @@ sm100_bf16_mega_moe_impl(void* y, break; } + // Read the source metadata of this warp's rows before the TMEM loads to overlap the latency + layout::TokenSrcMetadata cached_src_metadata[kNumRowsPerWarp]; + #pragma unroll + for (uint32_t j = 0; j < kNumRowsPerWarp; ++ j) { + const uint32_t m_idx_in_block = epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + j * 8 + warp_idx_in_wg * 2 + lane_idx / 16; + if (m_idx_in_block < valid_m) + cached_src_metadata[j] = task_info.is_shared() ? + layout::TokenSrcMetadata(sym_buffer.rank_idx, pool_m_idx + m_idx_in_block, kNumTopk) : + *workspace.get_token_src_metadata_ptr(pool_m_idx + m_idx_in_block); + } + __syncwarp(); + #pragma unroll for (uint32_t i = 0; i < STORE_BLOCK_M / ATOM_M; ++ i) { // Load from TMEM using .16x256b shape to satisfy STSM layout requirements @@ -1092,17 +1111,7 @@ sm100_bf16_mega_moe_impl(void* y, if (m_idx_in_block >= valid_m) break; - uint32_t dst_rank_idx, dst_token_idx, dst_topk_idx; - if (task_info.is_shared()) { - dst_rank_idx = sym_buffer.rank_idx; - dst_token_idx = pool_m_idx + m_idx_in_block; - dst_topk_idx = kNumTopk; - } else { - const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_m_idx + m_idx_in_block); - dst_rank_idx = src_metadata.rank_idx; - dst_token_idx = src_metadata.token_idx; - dst_topk_idx = src_metadata.topk_idx; - } + const auto& [dst_rank_idx, dst_token_idx, dst_topk_idx] = cached_src_metadata[j]; // Read from shared memory const auto smem_ptr = shared_storage.smem_d.l2[epilogue_wg_idx] + @@ -1131,19 +1140,8 @@ sm100_bf16_mega_moe_impl(void* y, if (epilogue_warp_idx == 0) Allocator().free(0, kNumTmemCols); - // NVLink barrier (grid sync + cross-rank signal + grid sync): ~4 us - comm::nvlink_barrier( - workspace, sym_buffer, sm_idx, epilogue_thread_idx, - [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } - ); - - // Barrier with dispatch warps, so that they can do clean workspace - ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); - // Combine: reduce top-k results and write back // NOTES: reuse shared memory from start up to the barriers - // 1 token, 1 topk latency: ~3 us constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); @@ -1180,97 +1178,124 @@ sm100_bf16_mega_moe_impl(void* y, return &shared_storage.combine_barriers[i + epilogue_warp_idx * 2]; }); - // Iterate over all tokens uint32_t combine_phase = 0; uint32_t load_stage_idx = 0; - for (uint32_t token_idx = sm_idx * kNumEpilogueWarps + epilogue_warp_idx; - token_idx < num_tokens; - token_idx += kNumSMs * kNumEpilogueWarps) { + + // Peers' grid indices for tagging their combine readiness; the load overlaps with the grid sync + DG_STATIC_ASSERT(kNumRanks <= kNumEpilogueThreads, "Insufficient threads for combine readiness"); + uint64_t peer_grid_idx = 0; + if (sm_idx == 0 and epilogue_thread_idx < kNumRanks) + peer_grid_idx = *workspace.get_peer_grid_idx_ptr(epilogue_thread_idx); + + // All local L2 writes are done after this grid sync + comm::grid_sync( + workspace, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + + // Notify remote ranks; ordered before the cleanup barrier by the dispatch/epilogue sync below + if (sm_idx == 0 and epilogue_thread_idx < kNumRanks) + ptx::st_rel_sys(sym_buffer.map(workspace.get_combine_ready_grid_idx_ptr(sym_buffer.rank_idx), epilogue_thread_idx), peer_grid_idx); + + // Barrier with dispatch warps, so that they can do clean workspace + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Iterate over all token chunks, 1 token 1 topk latency: ~3 us + const auto grid_idx = ptx::get_grid_idx() + 1; + for (uint32_t token_chunk_idx = epilogue_warp_idx * kNumSMs + sm_idx; token_chunk_idx < num_tokens * kNumChunks; token_chunk_idx += kNumSMs * kNumEpilogueWarps) { + const uint32_t token_idx = token_chunk_idx / kNumChunks; + const uint32_t chunk_idx = token_chunk_idx % kNumChunks; + // Read top-k slot indices: each lane reads one slot, then broadcast via exchange const int stored_topk_slot_idx = lane_idx < kNumTopk ? - static_cast(__ldg(buffer.input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : + static_cast(buffer.input_topk_idx_buffer.get_base_ptr()[token_idx * kNumTopk + lane_idx]) : (kNumSharedExperts > 0 and lane_idx == kNumTopk ? static_cast(kNumTopk) : -1); const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); - // Iterate all chunks - for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { - const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; - - // Move mask and load - uint32_t mask = total_mask; - const auto move_mask_and_load = [&](const uint32_t& i) { - if (mask) { - // Move - const uint32_t slot_idx = __ffs(mask) - 1; - mask ^= 1 << slot_idx; - - // Load - if (cute::elect_one_sync()) { - const auto src_ptr = math::advance_ptr( - buffer.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; - }; + // Wait for the ranks of the selected experts to finish their L2 writes + const bool is_routed = lane_idx < kNumTopk and stored_topk_slot_idx >= 0; + const auto peer_ready_ptr = workspace.get_combine_ready_grid_idx_ptr( + is_routed ? static_cast(stored_topk_slot_idx) / kNumExpertsPerRank : 0); + comm::wait_until([&]() { return __all_sync(0xffffffff, not is_routed or ptx::ld_acq_sys(peer_ready_ptr) == grid_idx); }, + [&]() { printf("DeepGEMM combine peers timeout: rank=%u, token=%u\n", sym_buffer.rank_idx, token_idx); }); - // Load the first selection - bool do_reduce = move_mask_and_load(load_stage_idx); + const uint32_t chunk_byte_offset = chunk_idx * kNumChunkBytes; - // 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); + // Move mask and load + uint32_t mask = total_mask; + const auto move_mask_and_load = [&](const uint32_t& i) { + if (mask) { + // Move + const uint32_t slot_idx = __ffs(mask) - 1; + mask ^= 1 << slot_idx; - // Accumulate - 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]); + // Load + if (cute::elect_one_sync()) { + const auto src_ptr = math::advance_ptr( + buffer.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); } - combine_phase ^= load_stage_idx; - load_stage_idx ^= 1; + __syncwarp(); + return true; } + return false; + }; + + // Load the first selection + bool do_reduce = 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); - // Cast + // Accumulate + combine_load_barriers[load_stage_idx]->wait(combine_phase); #pragma unroll for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { - uint4 casted; - auto casted_bf16 = reinterpret_cast(&casted); + 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) - casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); - - // Wait share memory release and write - 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::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); } - __syncwarp(); + cutlass::arch::fence_view_async_shared(); + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } - // TMA store the token chunk - 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(); + // Cast + #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]); + + // Wait share memory release and write + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); } - __syncwarp(); + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, + casted.x, casted.y, casted.z, casted.w); } + __syncwarp(); + + // TMA store the token chunk + 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 diff --git a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh index 0a626a14d4..7395c9c1a7 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh @@ -6,13 +6,18 @@ #include #include +#include #include +#include #include #include #include +#include #include -#include +#include +#include #include +#include namespace deep_gemm { @@ -22,17 +27,18 @@ template + typename epilogue_op_t> CUTLASS_GLOBAL void __launch_bounds__(kNumNonEpilogueThreads + kNumEpilogueThreads, 1) sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, + const __grid_constant__ epilogue_op_t epilogue_op, const __grid_constant__ cute::TmaDescriptor tensor_map_a, const __grid_constant__ cute::TmaDescriptor tensor_map_b, const __grid_constant__ cute::TmaDescriptor tensor_map_sfa, @@ -42,37 +48,65 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, using Barrier = cutlass::arch::ClusterTransactionBarrier; using Allocator = cute::conditional_t; - // C/D type: BF16 and FP32 are supported, with or without accumulation - DG_STATIC_ASSERT(cute::is_same_v or cute::is_same_v, "Invalid C/D data dtype"); + // NOTES: an FP4xFP4 pair uses the dedicated MXF4 MMA (packed 4-bit operands, K-major only, UMMA_K=64); + // all other combinations use the MXF8F6F4 MMA (unpacked operands, UMMA_K=32) + constexpr bool kIsMXF4 = cute::is_same_v and + cute::is_same_v; + + // The host launches the epilogue operator directly as a kernel argument + DG_STATIC_ASSERT(sizeof(epilogue_op_t) == sizeof(EpilogueArgs), + "Epilogue operators must not add state to `EpilogueArgs`"); + + // C/D type: BF16 and FP32 are supported, with or without accumulation; FP8 C/D requires the + // dynamically-scaled epilogue with per-32 UE8M0 SFD output, batched only, without accumulation + constexpr bool kWithOutputSF = cute::is_same_v; + DG_STATIC_ASSERT(kWithOutputSF ? (cute::is_same_v and + kGemmType == GemmType::Batched and not kWithAccumulation) : + (cute::is_same_v or cute::is_same_v), + "Invalid C/D data dtype"); // MMA Configs constexpr uint32_t LAYOUT_AD_M = 128; constexpr uint32_t UMMA_M = LAYOUT_AD_M * kNumMulticast; constexpr uint32_t UMMA_N = kSwapAB ? BLOCK_M : BLOCK_N; - constexpr uint32_t UMMA_K = 32; + constexpr uint32_t UMMA_K = kIsMXF4 ? 64 : 32; constexpr uint32_t LOAD_BLOCK_M = BLOCK_M / (kIsMulticastOnA ? kNumMulticast: 1); constexpr uint32_t LOAD_BLOCK_N = BLOCK_N / (kIsMulticastOnA ? 1 : kNumMulticast); - DG_STATIC_ASSERT(BLOCK_K == 128, "Invalid block K"); + constexpr uint32_t kPackFactor = get_smem_pack_factor(); + DG_STATIC_ASSERT(kPackFactor == get_smem_pack_factor(), "A/B SMEM pack factor must match"); + constexpr bool kIsFP4A = cute::is_same_v or + cute::is_same_v; + constexpr bool kIsFP4B = cute::is_same_v or + cute::is_same_v; + DG_STATIC_ASSERT((kIsMXF4 and BLOCK_K == 256) or (not kIsMXF4 and BLOCK_K == 128), "Invalid block K"); DG_STATIC_ASSERT(BLOCK_K % UMMA_K == 0, "Block K must be divisible by UMMA K"); DG_STATIC_ASSERT(kNumMulticast == 1 or kNumMulticast == 2, "Only support 1/2 multicast"); DG_STATIC_ASSERT((kSwapAB and BLOCK_N == LAYOUT_AD_M) or (not kSwapAB and (BLOCK_M == 32 or BLOCK_M == 64 or BLOCK_M == LAYOUT_AD_M)), "Invalid block size"); + // MXF4 only supports K-major operands and 32-element SF granularity + DG_STATIC_ASSERT(not kIsMXF4 or (kMajorA == cute::UMMA::Major::K and kMajorB == cute::UMMA::Major::K), "MXF4 only supports K-major"); + DG_STATIC_ASSERT(not kIsMXF4 or (kGranKA == 32 and kGranKB == 32), "MXF4 only supports kGranK=32"); + // SF configs constexpr uint32_t kNumUTCCPAlignedElems = 128; constexpr uint32_t SF_BLOCK_M = math::constexpr_align(BLOCK_M, kNumUTCCPAlignedElems); constexpr uint32_t SF_BLOCK_N = math::constexpr_align(BLOCK_N, kNumUTCCPAlignedElems); + constexpr uint32_t SF_BLOCK_K = BLOCK_K / 128; constexpr uint32_t kNumSFAStagesPerLoad = kGranKA == 32 ? 1 : 4; constexpr uint32_t kNumSFBStagesPerLoad = kGranKB == 32 ? 1 : 4; DG_STATIC_ASSERT(kGranKA == 32 or kGranKA == 128, "Invalid granularity K for A"); DG_STATIC_ASSERT(kGranKB == 32 or kGranKB == 128, "Invalid granularity K for B"); - DG_STATIC_ASSERT(not is_k_grouped_contiguous(kGemmType) or kGranKA == kGranKB, "K-grouped SF requires kGranKA == kGranKB"); - DG_STATIC_ASSERT(not is_k_grouped_contiguous(kGemmType) or kKAlignment % UMMA_K == 0, "K alignment must be divisible by UMMA K"); + DG_STATIC_ASSERT(not is_k_grouped_contiguous(kGemmType) or + (kGranKA == kGranKB and (kGranKA == 32 or kGranKA == 128)), + "K-grouped SF requires matching granularity K 32/128"); + DG_STATIC_ASSERT(not is_k_grouped_contiguous(kGemmType) or kKAlignment % BLOCK_K == 0, + "K alignment must be divisible by block K"); // Epilogue configs // Always enable pipeline for better performance constexpr uint32_t kNumEpilogueStages = 2; - constexpr uint32_t kNumTMAStoreStages = 2; + DG_STATIC_ASSERT(kNumTMAStoreStages == 1 or kNumTMAStoreStages == 2, "Invalid number of TMA store stages"); // NOTES: To maximize epilogue threads utilization, process an entire BLOCK_N // per store stage for swap-AB cases, and an entire BLOCK_M for non-swap cases constexpr uint32_t STORE_BLOCK_M = kSwapAB ? 16 : cute::min(BLOCK_M, LAYOUT_AD_M); @@ -80,26 +114,21 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, constexpr uint32_t kNumUMMAStoreThreads = kSwapAB ? kNumEpilogueThreads: STORE_BLOCK_M; DG_STATIC_ASSERT(kNumUMMAStoreThreads % 32 == 0, "Invalid store block M"); - // Share memory sizes - constexpr uint32_t SMEM_CD_SIZE_PER_STAGE = STORE_BLOCK_M * STORE_BLOCK_N * sizeof(cd_dtype_t); - constexpr uint32_t SMEM_CD_SIZE = SMEM_CD_SIZE_PER_STAGE * kNumTMAStoreStages; - constexpr uint32_t SMEM_A_SIZE_PER_STAGE = LOAD_BLOCK_M * BLOCK_K * sizeof(a_dtype_t); - constexpr uint32_t SMEM_B_SIZE_PER_STAGE = LOAD_BLOCK_N * BLOCK_K * sizeof(b_dtype_t); - constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = SF_BLOCK_M * sizeof(uint32_t); - constexpr uint32_t SMEM_SFB_SIZE_PER_STAGE = SF_BLOCK_N * sizeof(uint32_t); - DG_STATIC_ASSERT(SMEM_CD_SIZE % 1024 == 0 and SMEM_A_SIZE_PER_STAGE % 1024 == 0 and SMEM_B_SIZE_PER_STAGE % 1024 == 0, - "Shared memory of A/B must be aligned to 1024 bytes"); // NOTES: Make sure we have enough shared memory for UMMA padding - constexpr uint32_t UMMA_A_SIZE_PER_STAGE = math::constexpr_align(LOAD_BLOCK_M, LAYOUT_AD_M) * BLOCK_K * sizeof(a_dtype_t); - DG_STATIC_ASSERT(UMMA_A_SIZE_PER_STAGE <= SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE * kNumStages, "Memory Out of bound for UMMA"); + constexpr uint32_t UMMA_A_SIZE_PER_STAGE = math::constexpr_align(LOAD_BLOCK_M, LAYOUT_AD_M) * BLOCK_K * sizeof(a_dtype_t) / kPackFactor; // Tensor memory size and offsets constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumEpilogueStages; - constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M / 32; - constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N / 32; - constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); - constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols; - constexpr uint32_t kTmemStartColOfSFB = kNumAccumTmemCols + kNumSFATmemCols; + constexpr uint32_t kNumSFATmemCols = SF_BLOCK_M * SF_BLOCK_K / 32; + constexpr uint32_t kNumSFBTmemCols = SF_BLOCK_N * SF_BLOCK_K / 32; + constexpr uint32_t kNumSFTmemCols = kNumSFATmemCols + kNumSFBTmemCols; + constexpr uint32_t kNumOverlappedTmemCols = cute::max(kNumAccumTmemCols + kNumSFTmemCols, 512) - 512; + constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); + constexpr uint32_t kTmemStartColOfSFA = kNumAccumTmemCols - kNumOverlappedTmemCols; + constexpr uint32_t kTmemStartColOfSFB = kTmemStartColOfSFA + kNumSFATmemCols; + DG_STATIC_ASSERT(kNumOverlappedTmemCols <= UMMA_N, "Invalid overlapped tensor memory columns"); + DG_STATIC_ASSERT(kSwapAB or kNumOverlappedTmemCols == 0 or kNumOverlappedTmemCols <= STORE_BLOCK_N, + "Non-swap overlapped tensor memory columns must fit in the first epilogue store"); DG_STATIC_ASSERT(32 <= kNumTmemCols and kNumTmemCols <= 512, "Invalid tensor memory columns"); // Synchronize the cluster before 2-CTA TMEM allocation @@ -126,61 +155,39 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, const auto shape_sfa_k = math::ceil_div(shape_k, kGranKA * 4); const auto shape_sfb_k = math::ceil_div(shape_k, kGranKB * 4); - // Align to 1024 bytes for swizzle-128B - extern __shared__ __align__(1024) uint8_t smem_buffer[]; + using SharedStorage = layout::SM100FP8FP4GemmSharedStorage< + kNumStages, kNumEpilogueStages, kNumTMAStoreStages, LOAD_BLOCK_M, LOAD_BLOCK_N, BLOCK_K, + STORE_BLOCK_M, STORE_BLOCK_N, SF_BLOCK_M, SF_BLOCK_N, SF_BLOCK_K, + a_dtype_t, b_dtype_t, cd_dtype_t>; - // D/A/B shared memory - auto smem_cd = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + i * SMEM_CD_SIZE_PER_STAGE); - }); - auto smem_a = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + SMEM_CD_SIZE + i * SMEM_A_SIZE_PER_STAGE); - }); - auto smem_b = utils::PatternVisitor([&](const uint32_t& i) { - return reinterpret_cast(smem_buffer + SMEM_CD_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE); - }); - - // SFA/SFB shared memory - auto sf_start_ptr = reinterpret_cast(smem_b[kNumStages]); - auto smem_sfa = utils::PatternVisitor([=](const uint32_t& i) { - return reinterpret_cast(sf_start_ptr + i * SMEM_SFA_SIZE_PER_STAGE); - }); - auto smem_sfb = utils::PatternVisitor([=](const uint32_t& i) { - return reinterpret_cast(sf_start_ptr + kNumStages * SMEM_SFA_SIZE_PER_STAGE + i * SMEM_SFB_SIZE_PER_STAGE); - }); - - // Barriers and tensor memory pointer - auto barrier_start_ptr = reinterpret_cast(smem_sfb[kNumStages]);; - auto full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (i); }); - auto empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages + i); }); - auto with_sf_full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 2 + i); }); - auto tmem_full_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 3 + i); }); - auto tmem_empty_barriers = utils::PatternVisitor([=](const uint32_t& i) { return barrier_start_ptr + (kNumStages * 3 + kNumEpilogueStages + i); }); - auto tmem_ptr_in_smem = reinterpret_cast(barrier_start_ptr + kNumStages * 3 + kNumEpilogueStages * 2); + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + DG_STATIC_ASSERT(UMMA_A_SIZE_PER_STAGE <= sizeof(smem.a[0]) + sizeof(smem.b), "Memory out of bound for UMMA"); // Initialize barriers if (warp_idx == 1 and cute::elect_one_sync()) { #pragma unroll for (uint32_t i = 0; i < kNumStages; ++ i) { // Arrive at all CTAs - full_barriers[i]->init(1); - empty_barriers[i]->init(1); + smem.sf_full_barriers[i].init(1); + smem.empty_barriers[i].init(1); // Arrive only at the leader CTA - with_sf_full_barriers[i]->init(kNumMulticast * 32); + smem.full_barriers[i].init(kNumMulticast * (1 + 32 * SF_BLOCK_K)); } #pragma unroll for (uint32_t i = 0; i < kNumEpilogueStages; ++ i) { // Arrive at all CTAs - tmem_full_barriers[i]->init(1); + smem.tmem_full_barriers[i].init(1); // Arrive only at the leader CTA - tmem_empty_barriers[i]->init(kNumMulticast * kNumUMMAStoreThreads); + smem.tmem_empty_barriers[i].init(kNumMulticast * kNumUMMAStoreThreads); + smem.tmem_overlap_barriers[i].init(kNumMulticast * kNumUMMAStoreThreads); } // Make initialized barrier visible in async proxy cutlass::arch::fence_barrier_init(); } else if (warp_idx == 2) { // Allocate tensor memory - Allocator().allocate(kNumTmemCols, tmem_ptr_in_smem); + Allocator().allocate(kNumTmemCols, &smem.tmem_ptr); } kNumMulticast > 1 ? comm::cluster_sync_with_relaxed_arrive() : __syncthreads(); @@ -205,32 +212,40 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, // Dispatch warps into different roles if (warp_idx == 0 and cute::elect_one_sync()) { // TMA load warp + constexpr uint32_t kNumTMABytesPerStage = LOAD_BLOCK_M * BLOCK_K / (kIsFP4A ? 2 : 1) + + LOAD_BLOCK_N * BLOCK_K / (kIsFP4B ? 2 : 1); + // Persistently schedule over blocks while (scheduler.get_next_block(m_block_idx, n_block_idx)) { // Use dynamic load block M, when swap-AB is enabled const auto load_block_m = kSwapAB ? scheduler.get_aligned_effective_m_in_block(m_block_idx) / kNumMulticast : LOAD_BLOCK_M; // For k-grouped layout, the number of block K is variable - const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + const auto num_total_k_blocks = cute::max(1u, math::ceil_div(scheduler.current_shape_k, BLOCK_K)); for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { // Wait consumer release - empty_barriers[stage_idx]->wait(phase ^ 1); + smem.empty_barriers[stage_idx].wait(phase ^ 1); // Compute offsets - // NOTES: the group is always concatenated with the outer dimension - uint32_t m_idx = scheduler.template get_global_idx<(kGemmType == GemmType::MGroupedMasked), sched::IndexType::MN> ( + // M-grouped masked A stores groups along M; contiguous A uses global M. + constexpr bool kAWithGroupMNOffset = kGemmType == GemmType::MGroupedMasked; + uint32_t m_idx = scheduler.template get_global_idx ( shape_m, BLOCK_M, m_block_idx); - uint32_t n_idx = scheduler.template get_global_idx<(kMajorB == cute::UMMA::Major::K), sched::IndexType::MN> ( + // M-grouped K-major B stores groups along N. + constexpr bool kBWithGroupMNOffset = (kMajorB == cute::UMMA::Major::K) and + (is_m_grouped_contiguous(kGemmType) or kGemmType == GemmType::MGroupedMasked); + uint32_t n_idx = scheduler.template get_global_idx ( shape_n, BLOCK_N, n_block_idx, m_block_idx); - // NOTES: `k_idx` is actually the k index default for K-major, while `k_b_idx` may be MN-major - // And for all m-grouped GEMMs, A must be K-majored DG_STATIC_ASSERT(kGemmType == GemmType::Normal or is_k_grouped_contiguous(kGemmType) or kGemmType == GemmType::Batched or kMajorA == cute::UMMA::Major::K, "Invalid major"); - uint32_t k_idx = k_block_idx * BLOCK_K; - uint32_t k_a_idx = scheduler.template get_global_idx<(kMajorA == cute::UMMA::Major::MN), sched::IndexType::K> ( + // K-grouped A uses a shared K arena; other MN-major A groups along K. + constexpr bool kAWithGroupKOffset = is_k_grouped_contiguous(kGemmType) or (kMajorA == cute::UMMA::Major::MN); + // K-grouped B uses a shared K arena; M-grouped MN-major B groups along K. + constexpr bool kBWithGroupKOffset = is_k_grouped_contiguous(kGemmType) or (kMajorB == cute::UMMA::Major::MN); + uint32_t k_a_idx = scheduler.template get_global_idx ( shape_k, BLOCK_K, k_block_idx, m_block_idx); - uint32_t k_b_idx = scheduler.template get_global_idx<(kMajorB == cute::UMMA::Major::MN), sched::IndexType::K> ( + uint32_t k_b_idx = scheduler.template get_global_idx ( shape_k, BLOCK_K, k_block_idx, m_block_idx); // Add 2 CTA offsets @@ -239,43 +254,48 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, n_idx += kIsMulticastOnA ? 0 : (cute::block_rank_in_cluster() * LOAD_BLOCK_N); } - // Issue TMAs constexpr bool kIsBatchedMM = (kGemmType == GemmType::Batched); const uint32_t batch_idx = (kIsBatchedMM ? scheduler.current_group_idx : 0); - if constexpr (kMajorA == cute::UMMA::Major::K) - tma::copy( - &tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], k_a_idx, m_idx, 1, batch_idx); - if constexpr (kMajorA == cute::UMMA::Major::MN) - tma::copy( - &tensor_map_a, full_barriers[stage_idx], smem_a[stage_idx], m_idx, k_a_idx, 1, batch_idx); - if constexpr (kMajorB == cute::UMMA::Major::K) - tma::copy( - &tensor_map_b, full_barriers[stage_idx], smem_b[stage_idx], k_b_idx, n_idx, 1, batch_idx); - if constexpr (kMajorB == cute::UMMA::Major::MN) - tma::copy( - &tensor_map_b, full_barriers[stage_idx], smem_b[stage_idx], n_idx, k_b_idx, 1, batch_idx); - auto num_arrival_bytes = SMEM_A_SIZE_PER_STAGE / (std::is_same_v ? 1 : 2) + - SMEM_B_SIZE_PER_STAGE / (std::is_same_v ? 1 : 2); - // Issue SFA and SFB TMAs at certain stages - // No swizzling, so one TMA for one SF is enough + // Issue SFA and SFB TMAs first, so that the transpose can overlap the A/B transfer. + // No swizzling, so one TMA for one SF block is enough. SF loads land on their own barrier. + uint32_t sf_arrival_bytes = 0; if (k_block_idx % kNumSFAStagesPerLoad == 0) { uint32_t sfa_m_idx = m_block_idx * BLOCK_M; uint32_t sfa_k_idx = scheduler.template get_global_idx<(not is_m_grouped_contiguous(kGemmType)), sched::IndexType::SF_K>( - shape_sfa_k, 1, math::ceil_div(k_idx, BLOCK_K * kNumSFAStagesPerLoad)); - tma::copy(&tensor_map_sfa, full_barriers[stage_idx], smem_sfa[stage_idx], sfa_m_idx, sfa_k_idx); - num_arrival_bytes += BLOCK_M * sizeof(uint32_t); + shape_sfa_k, SF_BLOCK_K, k_block_idx / kNumSFAStagesPerLoad); + tma::copy(&tensor_map_sfa, &smem.sf_full_barriers[stage_idx], smem.sfa[stage_idx], sfa_m_idx, sfa_k_idx); + sf_arrival_bytes += sizeof(smem.sfa[0]); } if (k_block_idx % kNumSFBStagesPerLoad == 0) { uint32_t sfb_n_idx = n_block_idx * BLOCK_N; uint32_t sfb_k_idx = scheduler.template get_global_idx( - shape_sfb_k, 1, math::ceil_div(k_idx, BLOCK_K * kNumSFBStagesPerLoad), m_block_idx); - tma::copy(&tensor_map_sfb, full_barriers[stage_idx], smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx); - num_arrival_bytes += BLOCK_N * sizeof(uint32_t); + shape_sfb_k, SF_BLOCK_K, k_block_idx / kNumSFBStagesPerLoad, m_block_idx); + tma::copy(&tensor_map_sfb, &smem.sf_full_barriers[stage_idx], smem.sfb[stage_idx], sfb_n_idx, sfb_k_idx); + sf_arrival_bytes += sizeof(smem.sfb[0]); } + smem.sf_full_barriers[stage_idx].arrive_and_expect_tx(sf_arrival_bytes); + + // Issue A/B TMAs + if constexpr (kMajorA == cute::UMMA::Major::K) + tma::copy( + &tensor_map_a, &smem.full_barriers[stage_idx], smem.a[stage_idx], k_a_idx, m_idx, kNumMulticast, batch_idx); + if constexpr (kMajorA == cute::UMMA::Major::MN) + tma::copy( + &tensor_map_a, &smem.full_barriers[stage_idx], smem.a[stage_idx], m_idx, k_a_idx, kNumMulticast, batch_idx); + if constexpr (kMajorB == cute::UMMA::Major::K) + tma::copy( + &tensor_map_b, &smem.full_barriers[stage_idx], smem.b[stage_idx], k_b_idx, n_idx, kNumMulticast, batch_idx); + if constexpr (kMajorB == cute::UMMA::Major::MN) + tma::copy( + &tensor_map_b, &smem.full_barriers[stage_idx], smem.b[stage_idx], n_idx, k_b_idx, kNumMulticast, batch_idx); // Arrive at full barriers - full_barriers[stage_idx]->arrive_and_expect_tx(num_arrival_bytes); + if (is_leader_cta) { + smem.full_barriers[stage_idx].arrive_and_expect_tx(kNumTMABytesPerStage * kNumMulticast); + } else { + smem.full_barriers[stage_idx].arrive(0u); + } } } } else if (warp_idx == 1 and is_leader_cta) { @@ -289,10 +309,10 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, auto sf_desc = mma::sm100::make_sf_desc(nullptr); DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); - auto a_desc = mma::sm100::make_umma_desc(smem_a[0], 0, 0); - auto b_desc = mma::sm100::make_umma_desc(smem_b[0], 0, 0); - uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * SMEM_A_SIZE_PER_STAGE / 16 : 0u; - uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * SMEM_B_SIZE_PER_STAGE / 16 : 0u; + auto a_desc = mma::sm100::make_umma_desc(smem.a[0], 0, 0); + auto b_desc = mma::sm100::make_umma_desc(smem.b[0], 0, 0); + uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * sizeof(smem.a[0]) / 16 : 0u; + uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * sizeof(smem.b[0]) / 16 : 0u; // Checks for MMA instructions // NOTES: CUTLASS does not have such checks except the MMA traits, but we are not using these traits @@ -303,10 +323,12 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, // Persistently schedule over blocks while (scheduler.get_next_block(m_block_idx, n_block_idx)) { - // Wait tensor memory empty barrier arrival + // Wait until this accumulator stage is fully reusable. + // The overlap barrier is deferred to the first K block so that + // the SF UTCCP can be issued while the preceding epilogue drains. auto accum_stage_idx = scheduler.current_iter % kNumEpilogueStages; auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; - tmem_empty_barriers[accum_stage_idx]->wait(accum_phase_idx ^ 1); + smem.tmem_empty_barriers[accum_stage_idx].wait(accum_phase_idx ^ 1); ptx::tcgen05_after_thread_sync(); // Empty barrier arrival @@ -319,11 +341,11 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, cutlass::arch::umma_arrive_multicast_2x1SM(barrier, kCTAMask); } }; - umma_arrive(reinterpret_cast(empty_barriers[stage_idx])); + umma_arrive(reinterpret_cast(&smem.empty_barriers[stage_idx])); // NOTES: the tensor memory accumulator pipeline has nothing to do with multicasting if (do_tmem_full_arrive) - umma_arrive(reinterpret_cast(tmem_full_barriers[accum_stage_idx])); + umma_arrive(reinterpret_cast(&smem.tmem_full_barriers[accum_stage_idx])); __syncwarp(); }; @@ -334,85 +356,91 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, } // Launch MMAs - const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); - constexpr bool kMayHaveTailKBlock = is_k_grouped_contiguous(kGemmType) ? (kKAlignment % BLOCK_K != 0) : (SHAPE_K == 0 or SHAPE_K % BLOCK_K != 0); + const auto num_total_k_blocks = cute::max(1u, math::ceil_div(scheduler.current_shape_k, BLOCK_K)); #pragma unroll 4 for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { - // Wait TMA and SF-transpose arrival - with_sf_full_barriers[stage_idx]->wait(phase); - ptx::tcgen05_after_thread_sync(); - const auto a_desc_base_lo = ptx::exchange(a_desc_lo, stage_idx); const auto b_desc_base_lo = ptx::exchange(b_desc_lo, stage_idx); + + // Wait A/B TMA and SF-transpose arrival + smem.full_barriers[stage_idx].wait(phase); + ptx::tcgen05_after_thread_sync(); + + const uint32_t sfa_stage_in_group_idx = k_block_idx % kNumSFAStagesPerLoad; + const uint32_t sfb_stage_in_group_idx = k_block_idx % kNumSFBStagesPerLoad; if (cute::elect_one_sync()) { // Do SF copy at certain stages // TODO: process shared memory descriptor by addition using cute_utccp_t = cute::conditional_t; - const uint32_t sfa_stage_in_group_idx = k_block_idx % kNumSFAStagesPerLoad; if (sfa_stage_in_group_idx == 0) { #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { - auto smem_ptr = smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems; + for (uint32_t i = 0; i < SF_BLOCK_K * SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) { + auto smem_ptr = smem.sfa[stage_idx] + i * kNumUTCCPAlignedElems; mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); cute_utccp_t::copy(sf_desc, kTmemStartColOfSFA + i * 4); } } - const uint32_t sfb_stage_in_group_idx = k_block_idx % kNumSFBStagesPerLoad; if (sfb_stage_in_group_idx == 0) { #pragma unroll - for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { - auto smem_ptr = smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems; + for (uint32_t i = 0; i < SF_BLOCK_K * SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) { + auto smem_ptr = smem.sfb[stage_idx] + i * kNumUTCCPAlignedElems; mma::sm100::replace_smem_desc_addr(sf_desc, smem_ptr); cute_utccp_t::copy(sf_desc, kTmemStartColOfSFB + i * 4); } } + } + __syncwarp(); + + // The overlap barrier is waited after SF UTCCP but before the + // first UMMA: UTCCP can overlap the preceding epilogue, and + // subsequent K blocks share the same accumulator columns. + if constexpr (kNumOverlappedTmemCols > 0) { + if (k_block_idx == 0 and scheduler.current_iter > 0) { + ptx::tcgen05_before_thread_sync(); + const auto preceding_iter_idx = scheduler.current_iter - 1; + const auto preceding_stage_idx = preceding_iter_idx % kNumEpilogueStages; + const auto preceding_phase_idx = (preceding_iter_idx / kNumEpilogueStages) & 1; + smem.tmem_overlap_barriers[preceding_stage_idx].wait(preceding_phase_idx); + ptx::tcgen05_after_thread_sync(); + } + } + + if (cute::elect_one_sync()) { // Issue UMMA - using mma_t = cute::conditional_t< - kNumMulticast == 1, ptx::SM100_MMA_MXF8F6F4_SS, ptx::SM100_MMA_MXF8F6F4_2x1SM_SS>; - auto issue_umma = [&]() { - constexpr uint32_t kOffset = kUMMAKIdx * UMMA_K; - const uint32_t sfa_id = (kGranKA == 32 ? kUMMAKIdx : sfa_stage_in_group_idx); - const uint32_t sfb_id = (kGranKB == 32 ? kUMMAKIdx : sfb_stage_in_group_idx); + using mma_t = cute::conditional_t, + cute::conditional_t>; + #pragma unroll + for (uint32_t umma_k_idx = 0; umma_k_idx < BLOCK_K / UMMA_K; ++ umma_k_idx) { + const uint32_t offset = umma_k_idx * UMMA_K; + // Which 128-K SF sub-block this UMMA K step belongs to + const uint32_t subblock_idx = offset / kNumUTCCPAlignedElems; + // SF id (in units of 32 K-elements within the 128-K sub-block): + // gran-32 selects the 32-element SF at this UMMA K offset (MXF4's `scale_vec::2X` + // consumes two 32-element SFs per UMMA K=64 step, hence a step of 2); + // gran-128 uses the per-load stage id (whole 512-K SF span). + const uint32_t sf_id_in_subblock = (offset % kNumUTCCPAlignedElems) / 32; + const uint32_t tmem_col_sfa = kTmemStartColOfSFA + subblock_idx * SF_BLOCK_M / 32; + const uint32_t tmem_col_sfb = kTmemStartColOfSFB + subblock_idx * SF_BLOCK_N / 32; + const uint32_t sfa_id = (kGranKA == 32 ? sf_id_in_subblock : sfa_stage_in_group_idx); + const uint32_t sfb_id = (kGranKB == 32 ? sf_id_in_subblock : sfb_stage_in_group_idx); const auto runtime_instr_desc = kSwapAB ? mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sfb_id, sfa_id): mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sfa_id, sfb_id); - a_desc.lo = mma::sm100::advance_umma_desc_lo(a_desc_base_lo, 0, kOffset); - b_desc.lo = mma::sm100::advance_umma_desc_lo(b_desc_base_lo, 0, kOffset); + a_desc.lo = mma::sm100::advance_umma_desc_lo(a_desc_base_lo, 0, offset); + b_desc.lo = mma::sm100::advance_umma_desc_lo(b_desc_base_lo, 0, offset); if constexpr (kSwapAB) { - mma_t::fma(b_desc, a_desc, accum_stage_idx * UMMA_N, - kUMMAKIdx > 0 or k_block_idx > 0, runtime_instr_desc, - kTmemStartColOfSFB, kTmemStartColOfSFA); - } else { - mma_t::fma(a_desc, b_desc, accum_stage_idx * UMMA_N, - kUMMAKIdx > 0 or k_block_idx > 0, runtime_instr_desc, - kTmemStartColOfSFA, kTmemStartColOfSFB); - } - }; - auto issue_full_k_block = [&]() { - utils::for_each_static_until(std::make_integer_sequence(), issue_umma); - }; - - if constexpr (kMayHaveTailKBlock) { - auto issue_tail_k_block = [&](const uint32_t& remaining_k) { - const auto num_valid_umma_k = math::ceil_div(remaining_k, UMMA_K); - // Prefix expansion uses switch only for small cases to avoid long SASS. - utils::for_each_static_prefix(std::make_integer_sequence(), num_valid_umma_k, issue_umma); - }; - const auto is_last_k_block = k_block_idx == num_total_k_blocks - 1; - if (is_last_k_block) { - const auto remaining_k = scheduler.current_shape_k - k_block_idx * BLOCK_K; - if (remaining_k < BLOCK_K) - issue_tail_k_block(remaining_k); - else - issue_full_k_block(); + mma_t::fma(b_desc, a_desc, accum_stage_idx * (UMMA_N - kNumOverlappedTmemCols), + umma_k_idx > 0 or k_block_idx > 0, runtime_instr_desc, + tmem_col_sfb, tmem_col_sfa); } else { - issue_full_k_block(); + mma_t::fma(a_desc, b_desc, accum_stage_idx * (UMMA_N - kNumOverlappedTmemCols), + umma_k_idx > 0 or k_block_idx > 0, runtime_instr_desc, + tmem_col_sfa, tmem_col_sfb); } - } else { - issue_full_k_block(); } } __syncwarp(); @@ -427,10 +455,12 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, const auto iter_idx = scheduler.current_iter - 1; if (kNumMulticast > 1 and iter_idx >= 0) { const auto accum_phase_idx = (iter_idx / kNumEpilogueStages) & 1; - tmem_empty_barriers[iter_idx % kNumEpilogueStages]->wait(accum_phase_idx); + smem.tmem_empty_barriers[iter_idx % kNumEpilogueStages].wait(accum_phase_idx); } - } else if (warp_idx == 2) { + } else if (warp_idx == 2 or (SF_BLOCK_K == 2 and warp_idx == 3)) { // UTCCP transposer + // NOTES: use up to 2 warps to transpose, one per SF sub-block on K (only warp 3 when SF_BLOCK_K == 2) + const uint32_t sf_k_subblock_idx = warp_idx - 2; auto utccp_required_smem_warp_transpose = [&](const uint32_t* smem_ptr) { DG_STATIC_ASSERT(kNumUTCCPAlignedElems == 128, "Invalid aligned elements"); uint32_t values[4]; @@ -442,29 +472,26 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, }; while (scheduler.get_next_block(m_block_idx, n_block_idx)) { - const auto num_total_k_blocks = math::ceil_div(scheduler.current_shape_k, BLOCK_K); + const auto num_total_k_blocks = cute::max(1u, math::ceil_div(scheduler.current_shape_k, BLOCK_K)); for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) { - // Wait TMA arrival - full_barriers[stage_idx]->wait(phase); + // Wait SF TMA arrival + smem.sf_full_barriers[stage_idx].wait(phase); // Transpose for UTCCP at certain stages if (k_block_idx % kNumSFAStagesPerLoad == 0) { #pragma unroll for (uint32_t i = 0; i < SF_BLOCK_M / kNumUTCCPAlignedElems; ++ i) - utccp_required_smem_warp_transpose(smem_sfa[stage_idx] + i * kNumUTCCPAlignedElems); - // TODO: figure out whether the proxy fence is valid for 2-CTA cases - cutlass::arch::fence_view_async_shared(); + utccp_required_smem_warp_transpose(smem.sfa[stage_idx] + sf_k_subblock_idx * SF_BLOCK_M + i * kNumUTCCPAlignedElems); } if (k_block_idx % kNumSFBStagesPerLoad == 0) { #pragma unroll for (uint32_t i = 0; i < SF_BLOCK_N / kNumUTCCPAlignedElems; ++ i) - utccp_required_smem_warp_transpose(smem_sfb[stage_idx] + i * kNumUTCCPAlignedElems); - // TODO: figure out whether the proxy fence is valid for 2-CTA cases - cutlass::arch::fence_view_async_shared(); + utccp_required_smem_warp_transpose(smem.sfb[stage_idx] + sf_k_subblock_idx * SF_BLOCK_N + i * kNumUTCCPAlignedElems); } + // TODO: figure out whether the proxy fence is valid for 2-CTA cases + cutlass::arch::fence_view_async_shared(); - // Arrive - with_sf_full_barriers[stage_idx]->arrive(0u); + smem.full_barriers[stage_idx].arrive(0u); } } } else if (warp_idx >= kNumNonEpilogueThreads / 32 and warp_idx < (kNumNonEpilogueThreads + kNumUMMAStoreThreads) / 32) { @@ -474,7 +501,7 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, // NOTES: tensor memory addresses are simplified, as the hardware will ignore the warp index bits, // i.e., no need for `tmem_ptr |= (epilogue_warp_idx * 32) << 16`. // NOTES: we also forbid two CTAs to share the same SM and its tensor memory - DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(tmem_ptr_in_smem) == 0); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr) == 0); // Share store pipeline between blocks uint32_t tma_stage_idx = 0; @@ -485,36 +512,50 @@ sm100_fp8_fp4_gemm_1d1d_impl(int* grouped_layout, auto accum_phase_idx = (scheduler.current_iter / kNumEpilogueStages) & 1; // Wait UMMA arrival - tmem_full_barriers[accum_stage_idx]->wait(accum_phase_idx); + smem.tmem_full_barriers[accum_stage_idx].wait(accum_phase_idx); ptx::tcgen05_after_thread_sync(); - const auto tmem_base_addr = accum_stage_idx * UMMA_N; - const auto base_m_idx = scheduler.template get_global_idx<(not is_m_grouped_contiguous(kGemmType)), sched::IndexType::MN>(shape_m, BLOCK_M, m_block_idx); + const auto tmem_base_addr = accum_stage_idx * (UMMA_N - kNumOverlappedTmemCols); + const bool reverse_store_order = kNumOverlappedTmemCols > 0 and accum_stage_idx == 0; + DG_STATIC_ASSERT(kNumEpilogueStages == 2, + "reverse_store_order depends on exactly 2 accumulator stages"); + // Whether the group offset is encoded in the flattened CD M coordinate + constexpr bool kCDWithGroupOffset = not is_m_grouped_contiguous(kGemmType) and not is_k_grouped_contiguous(kGemmType); + const auto base_m_idx = scheduler.template get_global_idx(shape_m, BLOCK_M, m_block_idx); const auto base_n_idx = n_block_idx * BLOCK_N; + const bool is_empty_group = is_k_grouped_contiguous(kGemmType) and scheduler.current_shape_k == 0; if constexpr (kSwapAB) { const auto effective_m = scheduler.get_aligned_effective_m_in_block(m_block_idx); epilogue::sm100_store_cd_swap_ab< BLOCK_M, BLOCK_N, STORE_BLOCK_M, STORE_BLOCK_N, kSwizzleCDMode, kNumTMAStoreStages, kNumUMMAStoreThreads, - kGemmType, kWithAccumulation, - cd_dtype_t, epilogue_type_t> - (smem_cd, tma_stage_idx, tmem_base_addr, + kNumOverlappedTmemCols, + kGemmType, kWithAccumulation> + (smem, tma_stage_idx, tmem_base_addr, base_m_idx, base_n_idx, scheduler.current_group_idx, + is_empty_group, effective_m, epilogue_warp_idx, lane_idx, - tmem_empty_barriers[accum_stage_idx], + epilogue_op, + reverse_store_order, + &smem.tmem_overlap_barriers[accum_stage_idx], + &smem.tmem_empty_barriers[accum_stage_idx], tensor_map_cd); } else { epilogue::sm100_store_cd< BLOCK_M, BLOCK_N, STORE_BLOCK_M, STORE_BLOCK_N, kSwizzleCDMode, kNumTMAStoreStages, kNumUMMAStoreThreads, - kGemmType, kWithAccumulation, - cd_dtype_t, epilogue_type_t> - (smem_cd, tma_stage_idx, tmem_base_addr, + kNumOverlappedTmemCols, + kGemmType, kWithAccumulation> + (smem, tma_stage_idx, tmem_base_addr, base_m_idx, base_n_idx, scheduler.current_group_idx, + is_empty_group, epilogue_warp_idx, lane_idx, - tmem_empty_barriers[accum_stage_idx], + epilogue_op, + reverse_store_order, + &smem.tmem_overlap_barriers[accum_stage_idx], + &smem.tmem_empty_barriers[accum_stage_idx], tensor_map_cd); } } 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 4ac276d3ed..61fd17b070 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 @@ -35,6 +35,7 @@ template < uint32_t kNumSMs, uint32_t kNumRanks, float kActivationClamp, bool kFastMath, + typename weight_dtype_t, bool kHasShared = (kNumSharedExperts > 0), uint32_t L1_SHAPE_N = kIntermediateHidden * 2, uint32_t L1_SHAPE_K = kHidden, @@ -126,6 +127,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, ); const auto workspace = buffer.workspace; + using L2KBlockDependency = sched::L2KBlockDependency; + // SF and its buffer configs constexpr uint32_t kGranK = 32; constexpr uint32_t kNumUTCCPAlignedElems = 128; @@ -140,10 +143,11 @@ sm100_fp8_fp4_mega_moe_impl(void* y, }; // Data types - // NOTES: activations are FP8 (e4m3), weights are FP4 (e2m1) + // NOTES: activations and shared weights are FP8 (e4m3); routed weights may be FP8 or FP4 (e2m1) using a_dtype_t = cutlass::float_e4m3_t; - using b_dtype_t = cutlass::detail::float_e2m1_unpacksmem_t; using shared_b_dtype_t = cutlass::float_e4m3_t; + constexpr bool kIsWeightFP8 = cute::is_same_v; + DG_STATIC_ASSERT(kIsWeightFP8 or cute::is_same_v, "Invalid routed weight type"); // MMA configs // NOTES: always swap A/B, 2-CTA MMA, and matrices are K-major @@ -188,7 +192,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, alignas(kSharedMemoryAlignment) nv_bfloat16 l2[kNumEpilogueWarpgroups][STORE_BLOCK_M * BLOCK_N]; } smem_d; alignas(kSharedMemoryAlignment) a_dtype_t smem_a[kNumStages][LOAD_BLOCK_M * BLOCK_K]; - alignas(kSharedMemoryAlignment) b_dtype_t smem_b[kNumStages][LOAD_BLOCK_N * BLOCK_K]; + alignas(kSharedMemoryAlignment) weight_dtype_t smem_b[kNumStages][LOAD_BLOCK_N * BLOCK_K]; uint32_t smem_sfa[kNumStages][SF_BLOCK_M * (BLOCK_K / 128)]; uint32_t smem_sfb[kNumStages][SF_BLOCK_N * (BLOCK_K / 128)]; float2 amax_reduction[kNumEpilogueWarps][AMAX_REDUCTION_WARP_BUFFER_SIZE]; @@ -270,9 +274,10 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Allocate tensor memory Allocator().allocate(kNumTmemCols, &shared_storage.tmem_ptr_in_smem); } - // NOTES: Using `.relaxed` is allowed here since `fence_barrier_init` is `.release.cluster`, - // and `barrier.cluster.wait.aligned` is by default `.acquire` - comm::cluster_sync_with_relaxed_arrive(); + cute::cluster_sync(); + + // Wait for primary kernel completion + cudaGridDependencySynchronize(); // Task scheduler auto scheduler = sched::MegaMoEScheduler< @@ -307,8 +312,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // NVLink barrier tags constexpr uint32_t kBeforeDispatchPullBarrierTag = 1; - constexpr uint32_t kBeforeCombineReduceBarrierTag = 2; - constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 3; + constexpr uint32_t kAfterWorkspaceCleanBarrierTag = 2; // Adjust registers // NOTES: more experts per rank will cost more schedulers' registers @@ -344,7 +348,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, int expert_idx = -1; if (i + (lane_idx / kNumTopk) < num_tokens and lane_idx < kNumActivateLanes) { expert_idx = static_cast( - __ldg(buffer.input_topk_idx_buffer.get_base_ptr() + i * kNumTopk + lane_idx)); + buffer.input_topk_idx_buffer.get_base_ptr()[i * kNumTopk + lane_idx]); if (expert_idx >= 0) process(i * kNumTopk + lane_idx, expert_idx); } @@ -384,6 +388,12 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Write expert count if (sm_idx == 0) { + // Push this launch's grid index to every peer for tagging combine readiness; +1 differs from a zeroed workspace + DG_STATIC_ASSERT(kNumRanks <= kNumDispatchThreads, "Insufficient threads for the grid index push"); + if (thread_idx < kNumRanks) + *sym_buffer.map(workspace.get_peer_grid_idx_ptr(sym_buffer.rank_idx), thread_idx) = ptx::get_grid_idx() + 1; + __syncwarp(); + #pragma unroll for (uint32_t i = thread_idx; i < kNumExperts; i += kNumDispatchThreads) { const auto dst_rank_idx = i / kNumExpertsPerRank; @@ -556,6 +566,11 @@ sm100_fp8_fp4_mega_moe_impl(void* y, } __syncwarp(); + // Load the weight first, so that its remote latency overlaps with the SF copy below + const auto weight = *sym_buffer.map( + buffer.input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, + current_rank_in_expert_idx); + // Load and store SF (overlaps with last chunk's TMA load from remote) constexpr uint32_t kNumSFUint32 = kHidden / 128; DG_STATIC_ASSERT(kNumSFUint32 > 0 and kHidden % 128 == 0, "Invalid SF"); @@ -577,10 +592,6 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Store weights and metadata if (cute::elect_one_sync()) { - // Load weights - const auto weight = *sym_buffer.map( - buffer.input_topk_weights_buffer.get_base_ptr() + src_token_topk_idx, - current_rank_in_expert_idx); *buffer.l1_topk_weights_buffer.get_data_buffer(pool_token_idx % kNumRingTokens).template get_base_ptr() = weight; // Write source metadata for combine write-back (logical pool token) @@ -651,7 +662,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, for (uint32_t j = thread_idx; j < num_recv_m_blocks; j += kNumDispatchThreads) { *workspace.get_l1_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; *workspace.get_l1_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; - *workspace.get_l2_full_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; + *workspace.get_l2_full_mask_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; *workspace.get_l2_empty_count_ptr((expert_pool_block_offset + j) % kNumRingBlocks) = 0; } __syncwarp(); @@ -693,17 +704,17 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const auto ptr = workspace.get_l1_full_count_ptr(block_idx); const auto num_expected_tokens = BLOCK_M * (pool_block_idx / kNumRingBlocks + 1); while (ptx::ld_acq(ptr) != num_expected_tokens); - } else if (task_info.block_phase == sched::BlockPhase::Linear2) { - const auto ptr = workspace.get_l2_full_count_ptr(block_idx); - const auto num_expected_blocks = (L2_SHAPE_K / BLOCK_N) * 2 * (pool_block_idx / kNumRingBlocks + 1); - while (ptx::ld_acq(ptr) != num_expected_blocks); } else if (task_info.block_phase == sched::BlockPhase::SharedLinear2) { const auto ptr = workspace.get_shared_l2_full_count_ptr(block_idx); const auto num_expected_blocks = (SHARED_L2_SHAPE_K / BLOCK_N) * 2; while (ptx::ld_acq(ptr) != num_expected_blocks); } + L2KBlockDependency l2_k_block_dependency(workspace.get_l2_full_mask_ptr(ring_block_idx), pool_block_idx / kNumRingBlocks); for (uint32_t k_block_idx = 0; k_block_idx < num_k_blocks; advance_pipeline(k_block_idx)) { + if (task_info.block_phase == sched::BlockPhase::Linear2) + l2_k_block_dependency.wait(k_block_idx); + // Wait consumer release shared_storage.empty_barriers[stage_idx].wait(phase ^ 1); @@ -777,12 +788,14 @@ sm100_fp8_fp4_mega_moe_impl(void* y, shared_storage.full_barriers[stage_idx].arrive(0u); } } else { - tma::copy( + tma::copy( tensor_map_b_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_b[stage_idx], k_idx, n_idx, 2); tma::copy( tensor_map_sfb_ptr, &shared_storage.full_barriers[stage_idx], shared_storage.smem_sfb[stage_idx], sfb_n_idx, sfb_k_idx, 2); if (is_leader_cta) { - shared_storage.full_barriers[stage_idx].arrive_and_expect_tx(sizeof(SharedStorage::smem_b[0]) + sizeof(SharedStorage::smem_sfb[0]) * 2); + constexpr uint32_t kNumWeightBytes = sizeof(SharedStorage::smem_b[0]) * 2 / + (kIsWeightFP8 ? 1 : 2); + shared_storage.full_barriers[stage_idx].arrive_and_expect_tx(kNumWeightBytes + sizeof(SharedStorage::smem_sfb[0]) * 2); } else { shared_storage.full_barriers[stage_idx].arrive(0u); } @@ -800,7 +813,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Make instruction descriptor with block scaling // NOTES: always swap A/B auto routed_instr_desc = cute::UMMA::make_instr_desc_block_scaled< - b_dtype_t, a_dtype_t, float, cutlass::float_ue8m0_t, + weight_dtype_t, a_dtype_t, float, cutlass::float_ue8m0_t, UMMA_M, UMMA_N, cute::UMMA::Major::K, cute::UMMA::Major::K >(); @@ -811,13 +824,12 @@ sm100_fp8_fp4_mega_moe_impl(void* y, >(); auto sf_desc = mma::sm100::make_sf_desc(nullptr); + DG_STATIC_ASSERT(sizeof(weight_dtype_t) == sizeof(shared_b_dtype_t), "Weight SMEM descriptors must use identical addressing"); DG_STATIC_ASSERT(kNumStages <= 32, "Too many stages"); auto a_desc = mma::sm100::make_umma_desc(shared_storage.smem_a[0], 0, 0); auto b_desc = mma::sm100::make_umma_desc(shared_storage.smem_b[0], 0, 0); - auto shared_b_desc = mma::sm100::make_umma_desc(reinterpret_cast(shared_storage.smem_b[0]), 0, 0); - uint32_t a_desc_lo = lane_idx < kNumStages ? a_desc.lo + lane_idx * sizeof(SharedStorage::smem_a[0]) / 16 : 0u; - uint32_t b_desc_lo = lane_idx < kNumStages ? b_desc.lo + lane_idx * sizeof(SharedStorage::smem_b[0]) / 16 : 0u; - uint32_t shared_b_desc_lo = lane_idx < kNumStages ? shared_b_desc.lo + lane_idx * sizeof(SharedStorage::smem_b[0]) / 16 : 0u; + const uint32_t a_desc_lo = a_desc.lo; + const uint32_t b_desc_lo = b_desc.lo; // Checks for MMA instructions DG_STATIC_ASSERT((UMMA_M == 64 and UMMA_N % 8 == 0 and 8 <= UMMA_N and UMMA_N <= 256) or @@ -862,8 +874,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, shared_storage.full_barriers[stage_idx].wait(phase); ptx::tcgen05_after_thread_sync(); - const auto a_desc_base_lo = ptx::exchange(a_desc_lo, stage_idx); - const auto b_desc_base_lo = ptx::exchange(task_info.is_shared() ? shared_b_desc_lo : b_desc_lo, stage_idx); + const uint32_t a_desc_base_lo = a_desc_lo + stage_idx * sizeof(SharedStorage::smem_a[0]) / 16; + const uint32_t b_desc_base_lo = b_desc_lo + stage_idx * sizeof(SharedStorage::smem_b[0]) / 16; if (cute::elect_one_sync()) { #pragma unroll for (uint32_t umma_k_block_idx = 0; umma_k_block_idx < BLOCK_K / UMMA_BLOCK_K; ++ umma_k_block_idx) { @@ -889,13 +901,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, 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, umma_k_block_idx * UMMA_BLOCK_K * LOAD_BLOCK_M * sizeof(a_dtype_t), k * UMMA_K); - if (task_info.is_shared()) { - b_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, shared_b_dtype_t>(b_desc_base_lo, umma_k_block_idx * UMMA_BLOCK_K * LOAD_BLOCK_N * sizeof(shared_b_dtype_t), k * UMMA_K); - } else { - b_desc.lo = mma::sm100::advance_umma_desc_lo< - cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, b_dtype_t>(b_desc_base_lo, umma_k_block_idx * UMMA_BLOCK_K * LOAD_BLOCK_N * sizeof(b_dtype_t), k * UMMA_K); - } + b_desc.lo = mma::sm100::advance_umma_desc_lo< + cute::UMMA::Major::K, LOAD_BLOCK_N, kSwizzleBMode, weight_dtype_t>(b_desc_base_lo, umma_k_block_idx * UMMA_BLOCK_K * LOAD_BLOCK_N * sizeof(weight_dtype_t), k * UMMA_K); ptx::SM100_MMA_MXF8F6F4_2x1SM_SS::fma( b_desc, a_desc, accum_stage_idx * UMMA_N, k_block_idx > 0 or umma_k_block_idx > 0 or k > 0, runtime_instr_desc, @@ -1075,8 +1082,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, float2 thread_local_amax = {0.f, 0.f}; #pragma unroll for (uint32_t k = 0; k < 2; ++ k) { - thread_local_amax.x = cute::max(thread_local_amax.x, cute::abs(activation_values[i][k].x)); - thread_local_amax.y = cute::max(thread_local_amax.y, cute::abs(activation_values[i][k].y)); + thread_local_amax.x = fmaxf(thread_local_amax.x, fabsf(activation_values[i][k].x)); + thread_local_amax.y = fmaxf(thread_local_amax.y, fabsf(activation_values[i][k].y)); } // Amax reduction (warp-level) @@ -1103,12 +1110,14 @@ sm100_fp8_fp4_mega_moe_impl(void* y, // Reduce amax (warp-pair-level) const float2 wp_amax = shared_storage.amax_reduction[epilogue_warp_idx ^ 1][i * (ATOM_M / 2) + lane_idx % 4]; - 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); + amax_values[i].x = fmaxf(amax_values[i].x, wp_amax.x); + amax_values[i].y = fmaxf(amax_values[i].y, wp_amax.y); // Calculate SF - float2 sf, sf_inv; - math::get_e4m3_sf_and_sf_inv(amax_values[i], sf, sf_inv); + const uint2 sf_exp = {math::get_ue8m0_sf_exp(amax_values[i].x), + math::get_ue8m0_sf_exp(amax_values[i].y)}; + const float2 sf_inv = {math::get_ue8m0_sf_inv(sf_exp.x), + math::get_ue8m0_sf_inv(sf_exp.y)}; // Cast const float2 upper = __fmul2_rn(activation_values[i][0], sf_inv); @@ -1147,10 +1156,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, const auto sf_token_idx = block_idx * SF_BLOCK_M + transform_sf_token_idx(token_base_idx) + (lane_idx * 2) * 4; const auto sf_addr = k_uint_idx * mn_stride + sf_token_idx * static_cast(sizeof(uint32_t)) + byte_idx; - sf_base_ptr[sf_addr] = - (*reinterpret_cast(&sf.x) >> 23); - sf_base_ptr[sf_addr + 4 * static_cast(sizeof(uint32_t))] = - (*reinterpret_cast(&sf.y) >> 23); + sf_base_ptr[sf_addr] = static_cast(sf_exp.x); + sf_base_ptr[sf_addr + 4 * static_cast(sizeof(uint32_t))] = static_cast(sf_exp.y); } __syncwarp(); } @@ -1180,8 +1187,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, ptx::red_add_rel( workspace.get_shared_l2_full_count_ptr(pool_block_idx), 1u); } else { - ptx::red_add_rel( - workspace.get_l2_full_count_ptr(ring_block_idx), 1u); + L2KBlockDependency::arrive(workspace.get_l2_full_mask_ptr(ring_block_idx), n_block_idx); // Increment L1 empty count for this physical slot (one per N block) ptx::red_add( @@ -1213,6 +1219,18 @@ sm100_fp8_fp4_mega_moe_impl(void* y, break; } + // Read the source metadata of this warp's rows before the TMEM loads to overlap the latency + layout::TokenSrcMetadata cached_src_metadata[kNumRowsPerWarp]; + #pragma unroll + for (uint32_t j = 0; j < kNumRowsPerWarp; ++ j) { + const uint32_t m_idx_in_block = epilogue_wg_idx * WG_BLOCK_M + s * STORE_BLOCK_M + j * 8 + warp_idx_in_wg * 2 + lane_idx / 16; + if (m_idx_in_block < valid_m) + cached_src_metadata[j] = task_info.is_shared() ? + layout::TokenSrcMetadata(sym_buffer.rank_idx, pool_m_idx + m_idx_in_block, kNumTopk) : + *workspace.get_token_src_metadata_ptr(pool_m_idx + m_idx_in_block); + } + __syncwarp(); + #pragma unroll for (uint32_t i = 0; i < STORE_BLOCK_M / ATOM_M; ++ i) { // Load from TMEM using .16x256b shape to satisfy STSM layout requirements @@ -1271,17 +1289,7 @@ sm100_fp8_fp4_mega_moe_impl(void* y, if (m_idx_in_block >= valid_m) break; - uint32_t dst_rank_idx, dst_token_idx, dst_topk_idx; - if (task_info.is_shared()) { - dst_rank_idx = sym_buffer.rank_idx; - dst_token_idx = pool_m_idx + m_idx_in_block; - dst_topk_idx = kNumTopk; - } else { - const auto src_metadata = *workspace.get_token_src_metadata_ptr(pool_m_idx + m_idx_in_block); - dst_rank_idx = src_metadata.rank_idx; - dst_token_idx = src_metadata.token_idx; - dst_topk_idx = src_metadata.topk_idx; - } + const auto& [dst_rank_idx, dst_token_idx, dst_topk_idx] = cached_src_metadata[j]; // Read from shared memory const auto smem_ptr = reinterpret_cast(shared_storage.smem_d.l2[epilogue_wg_idx]) + @@ -1310,19 +1318,8 @@ sm100_fp8_fp4_mega_moe_impl(void* y, if (epilogue_warp_idx == 0) Allocator().free(0, kNumTmemCols); - // NVLink barrier (grid sync + cross-rank signal + grid sync): ~4 us - comm::nvlink_barrier( - workspace, sym_buffer, sm_idx, epilogue_thread_idx, - [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } - ); - - // Barrier with dispatch warps, so that they can do clean workspace - ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); - // Combine: reduce top-k results and write back // NOTES: reuse shared memory from start up to the barriers - // 1 token, 1 topk latency: ~3 us constexpr uint32_t kNumHiddenBytes = kHidden * sizeof(nv_bfloat16); constexpr uint32_t kNumElemsPerUint4 = sizeof(uint4) / sizeof(nv_bfloat162); @@ -1358,97 +1355,124 @@ sm100_fp8_fp4_mega_moe_impl(void* y, return &shared_storage.combine_barriers[i + epilogue_warp_idx * 2]; }); - // Iterate over all tokens uint32_t combine_phase = 0; uint32_t load_stage_idx = 0; - for (uint32_t token_idx = sm_idx * kNumEpilogueWarps + epilogue_warp_idx; - token_idx < num_tokens; - token_idx += kNumSMs * kNumEpilogueWarps) { + + // Peers' grid indices for tagging their combine readiness; the load overlaps with the grid sync + DG_STATIC_ASSERT(kNumRanks <= kNumEpilogueThreads, "Insufficient threads for combine readiness"); + uint64_t peer_grid_idx = 0; + if (sm_idx == 0 and epilogue_thread_idx < kNumRanks) + peer_grid_idx = *workspace.get_peer_grid_idx_ptr(epilogue_thread_idx); + + // All local L2 writes are done after this grid sync + comm::grid_sync( + workspace, sm_idx, epilogue_thread_idx, + [&]() { ptx::sync_aligned(kNumEpilogueThreads, kEpilogueFullBarrierIdx); } + ); + + // Notify remote ranks; ordered before the cleanup barrier by the dispatch/epilogue sync below + if (sm_idx == 0 and epilogue_thread_idx < kNumRanks) + ptx::st_rel_sys(sym_buffer.map(workspace.get_combine_ready_grid_idx_ptr(sym_buffer.rank_idx), epilogue_thread_idx), peer_grid_idx); + + // Barrier with dispatch warps, so that they can do clean workspace + ptx::sync_unaligned(kNumDispatchThreads + kNumEpilogueThreads, kDispatchWithEpilogueBarrierIdx); + + // Iterate over all token chunks, 1 token 1 topk latency: ~3 us + const auto grid_idx = ptx::get_grid_idx() + 1; + for (uint32_t token_chunk_idx = epilogue_warp_idx * kNumSMs + sm_idx; token_chunk_idx < num_tokens * kNumChunks; token_chunk_idx += kNumSMs * kNumEpilogueWarps) { + const uint32_t token_idx = token_chunk_idx / kNumChunks; + const uint32_t chunk_idx = token_chunk_idx % kNumChunks; + // Read top-k slot indices: each lane reads one slot, then broadcast via exchange const int stored_topk_slot_idx = lane_idx < kNumTopk ? - static_cast(__ldg(buffer.input_topk_idx_buffer.get_base_ptr() + token_idx * kNumTopk + lane_idx)) : + static_cast(buffer.input_topk_idx_buffer.get_base_ptr()[token_idx * kNumTopk + lane_idx]) : (kNumSharedExperts > 0 and lane_idx == kNumTopk ? static_cast(kNumTopk) : -1); const uint32_t total_mask = __ballot_sync(0xffffffff, stored_topk_slot_idx >= 0); - // Iterate all chunks - for (uint32_t chunk = 0; chunk < kNumChunks; ++ chunk) { - const uint32_t chunk_byte_offset = chunk * kNumChunkBytes; - - // Move mask and load - uint32_t mask = total_mask; - const auto move_mask_and_load = [&](const uint32_t& i) { - if (mask) { - // Move - const uint32_t slot_idx = __ffs(mask) - 1; - mask ^= 1 << slot_idx; - - // Load - if (cute::elect_one_sync()) { - const auto src_ptr = math::advance_ptr( - buffer.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; - }; + // Wait for the ranks of the selected experts to finish their L2 writes + const bool is_routed = lane_idx < kNumTopk and stored_topk_slot_idx >= 0; + const auto peer_ready_ptr = workspace.get_combine_ready_grid_idx_ptr( + is_routed ? static_cast(stored_topk_slot_idx) / kNumExpertsPerRank : 0); + comm::wait_until([&]() { return __all_sync(0xffffffff, not is_routed or ptx::ld_acq_sys(peer_ready_ptr) == grid_idx); }, + [&]() { printf("DeepGEMM combine peers timeout: rank=%u, token=%u\n", sym_buffer.rank_idx, token_idx); }); - // Load the first selection - bool do_reduce = move_mask_and_load(load_stage_idx); + const uint32_t chunk_byte_offset = chunk_idx * kNumChunkBytes; - // 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); + // Move mask and load + uint32_t mask = total_mask; + const auto move_mask_and_load = [&](const uint32_t& i) { + if (mask) { + // Move + const uint32_t slot_idx = __ffs(mask) - 1; + mask ^= 1 << slot_idx; - // Accumulate - 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]); + // Load + if (cute::elect_one_sync()) { + const auto src_ptr = math::advance_ptr( + buffer.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); } - combine_phase ^= load_stage_idx; - load_stage_idx ^= 1; + __syncwarp(); + return true; } + return false; + }; + + // Load the first selection + bool do_reduce = move_mask_and_load(load_stage_idx); - // Cast + // 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); + + // Accumulate + combine_load_barriers[load_stage_idx]->wait(combine_phase); #pragma unroll for (uint32_t j = 0; j < kNumUint4PerLane; ++ j) { - uint4 casted; - auto casted_bf16 = reinterpret_cast(&casted); + 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) - casted_bf16[l] = __float22bfloat162_rn(reduced[j * kNumElemsPerUint4 + l]); - - // Wait share memory release and write - 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::accumulate(reduced[j * kNumElemsPerUint4 + l], bf16_values[l]); } - __syncwarp(); + cutlass::arch::fence_view_async_shared(); + combine_phase ^= load_stage_idx; + load_stage_idx ^= 1; + } - // TMA store the token chunk - 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(); + // Cast + #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]); + + // Wait share memory release and write + if (j == 0) { + ptx::tma_store_wait<0>(); + __syncwarp(); } - __syncwarp(); + ptx::st_shared(combine_store_buffer + j * 32 + lane_idx, + casted.x, casted.y, casted.z, casted.w); } + __syncwarp(); + + // TMA store the token chunk + 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 diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mega_mhc.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mega_mhc.cuh new file mode 100644 index 0000000000..eaf618c1ba --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_mega_mhc.cuh @@ -0,0 +1,704 @@ +#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 + +namespace deep_gemm { + +using namespace layout::mega_mhc; + +namespace mega_mhc { + +// A task is one (m-block, K-split) pair, subdivided into BLOCK_K blocks. +// UMMA_K is the per-route hidden width of one TF32 UMMA instruction. +constexpr uint32_t UMMA_K = 32 / sizeof(float); +constexpr uint32_t kNumTokenRowsPerLane = 2; + +} // namespace mega_mhc + +// Persistent mHC with static grid-stride roles and TMA-native layouts. +// Norm receives one arrival only after Shifted X1 stores and statistics are committed. +template +CUTLASS_GLOBAL void __launch_bounds__(kNumThreads, 1) + sm100_mega_mhc_impl(const __grid_constant__ cute::TmaDescriptor tensor_map_residual, + const __grid_constant__ cute::TmaDescriptor tensor_map_x, + const __grid_constant__ cute::TmaDescriptor tensor_map_fn, + const __grid_constant__ cute::TmaDescriptor tensor_map_post_mix, + const __grid_constant__ cute::TmaDescriptor tensor_map_comb_res_mix, + const __grid_constant__ cute::TmaDescriptor tensor_map_shifted_prev_mix, + const __grid_constant__ cute::TmaDescriptor tensor_map_new_residual, + const __grid_constant__ cute::TmaDescriptor tensor_map_y_bf16, + const __grid_constant__ MixArgs mix_args, + const __grid_constant__ NormArgs norm_args, + void* gmem_scratch, + uint64_t* gmem_split_barriers) { +#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000)) or defined(__CLION_IDE__) + using namespace mega_mhc; + using Allocator = cute::TMEM::Allocator1Sm; + // Each Post worker occupies exactly one WG; the role dispatch hardwires two of them. + constexpr uint32_t kNumPostWGs = 2; + + // Template checks + DG_STATIC_ASSERT(kHidden > 0 and kHidden % BLOCK_K == 0 and kNumSplits <= kHidden / BLOCK_K, + "Splits must contain nonempty, aligned K blocks"); + DG_STATIC_ASSERT(SF_BLOCK_M == 0 or kStoreFP8, "Shared SF requires FP8 output"); + + // MMA and tensor memory configs + constexpr uint32_t kNumTmemColumnsPerAStage = kNumRoutes * UMMA_K; + constexpr uint32_t kAccumTmemStartColumn = kNumAStages * kNumTmemColumnsPerAStage; + constexpr uint32_t kNumTmemCols = + utils::get_num_aligned_tmem_cols(); + // Task geometry + constexpr uint32_t kNumKBlocksPerSplit = kHidden / BLOCK_K / kNumSplits; + constexpr uint32_t kNumLongSplits = kHidden / BLOCK_K % kNumSplits; + const auto get_task_k_begin = [&](const uint32_t split_idx) { + return split_idx * (kNumKBlocksPerSplit * BLOCK_K) + cute::min(split_idx, kNumLongSplits) * BLOCK_K; + }; + const auto get_num_task_k_blocks = [&](const uint32_t split_idx) { + return kNumKBlocksPerSplit + (split_idx < kNumLongSplits); + }; + + // CTA role assignments + constexpr uint32_t kFirstPostWGIdx = 1; + constexpr uint32_t kWorkspaceEpilogueWGIdx = kFirstPostWGIdx + kNumPostWGs; + constexpr uint32_t kMixWGIdx = kWorkspaceEpilogueWGIdx + 1; + constexpr uint32_t kNormWGIdx = kMixWGIdx + 1; + constexpr uint32_t kNumWGsPerCTA = kNormWGIdx + 1; + constexpr uint32_t kNumPostWarps = kNumPostWGs * kNumWarpsPerWG; + DG_STATIC_ASSERT(kNumWGsPerCTA == 6, "Unexpected CTA topology"); + + // Register reconfigurations + constexpr uint32_t kNumControlRegisters = 32; + constexpr uint32_t kNumPostRegisters = kIsShifted ? 136 : 112; + constexpr uint32_t kNumLightweightRegisters = 40; + constexpr uint32_t kNumShiftedNormRegisters = 96; + constexpr uint32_t kNumNormalRegisters = 80; + + // WG3 epilogue and WG4 Mix use the same lightweight allocation. + constexpr uint32_t kStaticRoleRegisterBudget = kNumControlRegisters + kNumPostRegisters * kNumPostWGs + + 2 * kNumLightweightRegisters + kNumShiftedNormRegisters; + constexpr uint32_t kMaxRegisterBudget = cute::max(kStaticRoleRegisterBudget, kNumNormalRegisters * kNumWGsPerCTA); + DG_STATIC_ASSERT(kMaxRegisterBudget * 128 <= 64512, "Register reconfiguration exceeds the SM budget"); + DG_STATIC_ASSERT(kMaxRegisterBudget <= 80 * kNumWGsPerCTA, "Register reconfiguration exceeds the CTA entry pool"); + + // Thread indices + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane_idx = ptx::get_lane_idx(); + + // Post and workspace-epilogue warps share this mapping: each lane owns two rows and one column pair. + // NOTES: runtime indices are recomputed per role branch so no value stays live across the + // register reconfigurations (which would spill to local memory). + constexpr uint32_t kNumRowsPerWarp = BLOCK_M / kNumWarpsPerWG; + constexpr uint32_t kNumRowGroupsPerWarp = kNumRowsPerWarp / kNumTokenRowsPerLane; + constexpr uint32_t kNumLanesPerRowGroup = 32 / kNumRowGroupsPerWarp; + const auto get_first_row_idx = [&] { + return warp_idx % kNumWarpsPerWG * kNumRowsPerWarp + lane_idx / kNumLanesPerRowGroup; + }; + + // Runtime state + const uint32_t num_m_blocks = Workspace::get_num_m_blocks(norm_args.num_tokens); + const uint32_t num_tasks = num_m_blocks * kNumSplits; + Workspace workspace(gmem_scratch, num_m_blocks); + + // Shared memory + extern __shared__ __align__(kSwizzleAlignment) uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + + const auto init_barriers = [](auto& barriers, const uint32_t count) { + #pragma unroll + for (auto& barrier: barriers) + barrier.init(count); + }; + + // Descriptor cache hints and barriers are initialized once per persistent CTA. + if (warp_idx == 0 and cute::elect_one_sync()) { + smem.next_norm_task_ticket = 0; + + // Normal aliases the shifted-only descriptors to valid maps, so prefetch unconditionally. + for (const auto& tensor_map: {&tensor_map_residual, &tensor_map_x, &tensor_map_fn, + &tensor_map_post_mix, &tensor_map_comb_res_mix, + &tensor_map_shifted_prev_mix, &tensor_map_new_residual, + &tensor_map_y_bf16}) + cute::prefetch_tma_descriptor(tensor_map); + + init_barriers(smem.full_io_barriers, 1); + init_barriers(smem.empty_io_barriers, 1); + init_barriers(smem.store_ready_io_barriers, kNumPostWarps); + init_barriers(smem.full_coeff_barriers, 1); + init_barriers(smem.empty_coeff_barriers, kNumPostWarps); + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 1 and cute::elect_one_sync()) { + // Fn load pipeline. + init_barriers(smem.full_fn_barriers, 1); + init_barriers(smem.empty_fn_barriers, 1); + cutlass::arch::fence_barrier_init(); + } else if (warp_idx == 2 and cute::elect_one_sync()) { + // MMA, Post, and workspace-epilogue pipelines. + init_barriers(smem.full_a_barriers, 128); + init_barriers(smem.empty_a_barriers, 1); + init_barriers(smem.full_accum_barriers, 1); + init_barriers(smem.empty_accum_barriers, kNumWarpsPerWG); + + smem.full_stats_barrier.init(kNumPostWarps); + smem.empty_stats_barrier.init(kNumWarpsPerWG + (kIsShifted ? 1 : 0)); + smem.mix_arrival_barrier.init(kNumWarpsPerWG + (kIsShifted ? 0 : 1)); + cutlass::arch::fence_barrier_init(); + } + + // Warp 3 owns the CTA-wide TMEM allocation. + if (warp_idx == 3) + Allocator().allocate(kNumTmemCols, &smem.tmem_base); + + // Local initialization can overlap the preceding grid, but persistent split barriers cannot. + cudaGridDependencySynchronize(); + if (warp_idx == kFirstPostWGIdx * kNumWarpsPerWG + 1 and cute::elect_one_sync()) { + for (uint32_t m_block_idx = blockIdx.x; m_block_idx < num_m_blocks; m_block_idx += kNumSMs) + sched::mega_mhc::Mix::init(gmem_split_barriers, m_block_idx); + } + if constexpr (kIsShifted) { + if (warp_idx == kFirstPostWGIdx * kNumWarpsPerWG and cute::elect_one_sync()) { + for (uint32_t m_block_idx = blockIdx.x; m_block_idx < num_m_blocks; m_block_idx += kNumSMs) + sched::mega_mhc::Norm::init(gmem_split_barriers, m_block_idx); + } + } + + __syncthreads(); + const uint32_t tmem_base = ptx::ld_shared(&smem.tmem_base); + + // Static roles follow the execution graph: WG0 TMA/MMA, WG1-2 Post, WG3 workspace epilogue, WG4 Mix, and WG5 Norm. + if (warp_idx < kFirstPostWGIdx * kNumWarpsPerWG) { + // WG0: producer pipelines + cutlass::arch::warpgroup_reg_dealloc(); + + // IO and Fn warp workers follow the same task traversal; only their hidden-block copies differ. + const auto run_tma_load_worker = [&](auto& empty_barriers, const auto& load_hidden_block) { + constexpr uint32_t kNumStages = sizeof(empty_barriers) / sizeof(empty_barriers[0]); + RingPipeline pipeline; + + if (cute::elect_one_sync()) { + for (uint32_t task_idx = blockIdx.x; task_idx < num_tasks; task_idx += kNumSMs) { + const uint32_t split_idx = task_idx % kNumSplits; + const uint32_t num_task_k_blocks = get_num_task_k_blocks(split_idx); + const uint32_t m_idx = task_idx / kNumSplits * BLOCK_M; + + #pragma unroll 1 + for (uint32_t k_block_idx = 0; k_block_idx < num_task_k_blocks; ++ k_block_idx) { + CUTE_TIE_DECL(pipeline.advance(), stage_idx, stage_phase); + empty_barriers[stage_idx].wait(stage_phase ^ 1u); + load_hidden_block(stage_idx, get_task_k_begin(split_idx) + k_block_idx * BLOCK_K, m_idx); + } + } + } + __syncwarp(); + }; + + if (warp_idx == 0) { + // Warp 0: residual and X TMA-load worker + run_tma_load_worker(smem.empty_io_barriers, + [&](const uint32_t& stage_idx, const uint32_t& k_idx, const uint32_t& m_idx) { + tma::copy( + &tensor_map_residual, &smem.full_io_barriers[stage_idx], + reinterpret_cast(smem.residual[stage_idx][0]), k_idx, m_idx); + tma::copy( + &tensor_map_x, &smem.full_io_barriers[stage_idx], + reinterpret_cast(smem.x[stage_idx]), k_idx, m_idx); + smem.full_io_barriers[stage_idx].arrive_and_expect_tx(sizeof(smem.residual[0]) + sizeof(smem.x[0])); + }); + } else if (warp_idx == 1) { + // Warp 1: Fn TMA-load worker + run_tma_load_worker(smem.empty_fn_barriers, + [&](const uint32_t& stage_idx, const uint32_t& k_idx, const uint32_t&) { + tma::copy( + &tensor_map_fn, &smem.full_fn_barriers[stage_idx], + smem.fn[stage_idx][0][0], k_idx, 0u); + smem.full_fn_barriers[stage_idx].arrive_and_expect_tx(sizeof(smem.fn[0])); + }); + } else if (warp_idx == 2) { + // Warp 2: HC MMA worker + using umma_t = cute::SM100_MMA_TF32_TS; + const auto runtime_instr_desc = cute::UMMA::make_runtime_instr_desc(); + + RingPipeline fn_pipeline; + RingPipeline a_pipeline; + RingPipeline accum_pipeline; + + for (uint32_t task_idx = blockIdx.x; task_idx < num_tasks; task_idx += kNumSMs) { + const uint32_t split_idx = task_idx % kNumSplits; + const uint32_t num_task_k_blocks = get_num_task_k_blocks(split_idx); + + // Wait workspace-epilogue accumulator release + CUTE_TIE_DECL(accum_pipeline.advance(), accum_stage_idx, accum_phase); + smem.empty_accum_barriers[accum_stage_idx].wait(accum_phase ^ 1u); + ptx::tcgen05_after_thread_sync(); + + #pragma unroll 1 + for (uint32_t k_block_idx = 0; k_block_idx < num_task_k_blocks; ++ k_block_idx) { + // Wait TMA Fn arrival + CUTE_TIE_DECL(fn_pipeline.advance(), fn_stage_idx, fn_phase); + smem.full_fn_barriers[fn_stage_idx].wait(fn_phase); + ptx::tcgen05_after_thread_sync(); + + #pragma unroll + for (uint32_t umma_k_idx = 0; umma_k_idx < BLOCK_K / UMMA_K; ++ umma_k_idx) { + // Wait Post A arrival + CUTE_TIE_DECL(a_pipeline.advance(), a_stage_idx, a_phase); + smem.full_a_barriers[a_stage_idx].wait(a_phase); + ptx::tcgen05_after_thread_sync(); + + // Hidden blocks, UMMA K atoms, and routes accumulate into one HC output block in fixed order. + #pragma unroll + for (uint32_t route_idx = 0; route_idx < kNumRoutes; ++ route_idx) { + auto b_desc = mma::sm100::make_umma_desc( + smem.fn[fn_stage_idx][umma_k_idx * UMMA_K / kFnAtomK][route_idx], + 0, umma_k_idx * UMMA_K % kFnAtomK); + umma_t::fma(tmem_base + a_stage_idx * kNumTmemColumnsPerAStage + route_idx * UMMA_K, + b_desc, tmem_base + kAccumTmemStartColumn + accum_stage_idx * BLOCK_N, + k_block_idx > 0 or umma_k_idx > 0 or route_idx > 0, runtime_instr_desc); + } + // Release the A stage + cutlass::arch::umma_arrive(reinterpret_cast(&smem.empty_a_barriers[a_stage_idx])); + } + // Release the Fn stage + cutlass::arch::umma_arrive(reinterpret_cast(&smem.empty_fn_barriers[fn_stage_idx])); + } + // Commit this split's accumulator + cutlass::arch::umma_arrive(reinterpret_cast(&smem.full_accum_barriers[accum_stage_idx])); + } + } else { + // Warp 3: output-store, coefficient-prefetch, and Norm-arrival worker + RingPipeline io_pipeline; + RingPipeline coeff_pipeline; + + const auto issue_coeff_tma = [&](const uint32_t& m_idx, const uint32_t& stage_idx) { + tma::copy(&tensor_map_post_mix, &smem.full_coeff_barriers[stage_idx], + smem.post_coeff[stage_idx][0], 0u, m_idx); + tma::copy( + &tensor_map_comb_res_mix, &smem.full_coeff_barriers[stage_idx], smem.comb_coeff[stage_idx][0], 0u, m_idx); + if constexpr (kIsShifted) + tma::copy(&tensor_map_shifted_prev_mix, &smem.full_coeff_barriers[stage_idx], + smem.pre_coeff[stage_idx][0], 0u, m_idx); + smem.full_coeff_barriers[stage_idx].arrive_and_expect_tx(sizeof(smem.post_coeff[0]) + sizeof(smem.comb_coeff[0]) + + (kIsShifted ? sizeof(smem.pre_coeff[0]) : 0)); + }; + + // The coefficient ring advances once per task, independently of the per-block IO ring. + auto advance_coeff_pipeline = [&]() { + CUTE_TIE_DECL(coeff_pipeline.advance(), stage_idx, phase); + smem.empty_coeff_barriers[stage_idx].wait(phase ^ 1u); + return stage_idx; + }; + + // The first task has no preceding store-drain window, so seed its coefficients here. + if (blockIdx.x < num_tasks) { + const uint32_t coeff_stage_idx = advance_coeff_pipeline(); + if (cute::elect_one_sync()) + issue_coeff_tma(blockIdx.x / kNumSplits * BLOCK_M, coeff_stage_idx); + __syncwarp(); + } + + for (uint32_t task_idx = blockIdx.x, task_phase = 0; task_idx < num_tasks; task_idx += kNumSMs, task_phase ^= 1u) { + const uint32_t split_idx = task_idx % kNumSplits; + const uint32_t num_task_k_blocks = get_num_task_k_blocks(split_idx); + const uint32_t m_block_idx = task_idx / kNumSplits; + if constexpr (kIsShifted) + sched::mega_mhc::Norm::wait_init(gmem_split_barriers, m_block_idx); + uint32_t next_coeff_stage_idx = 0; + + #pragma unroll 1 + for (uint32_t k_block_idx = 0; k_block_idx < num_task_k_blocks; ++ k_block_idx) { + // Wait for all Post route stores of this hidden block + CUTE_TIE_DECL(io_pipeline.advance(), io_stage_idx, io_phase); + smem.store_ready_io_barriers[io_stage_idx].wait(io_phase); + + const uint32_t k_idx = get_task_k_begin(split_idx) + k_block_idx * BLOCK_K; + const uint32_t m_idx = m_block_idx * BLOCK_M; + const bool should_prefetch_next_coeff = + k_block_idx == num_task_k_blocks - 1 and task_idx + kNumSMs < num_tasks; + + // Wait for a reusable slot before entering the final store-drain overlap window. + if (should_prefetch_next_coeff) + next_coeff_stage_idx = advance_coeff_pipeline(); + + if (cute::elect_one_sync()) { + // New residual is re-read by the fused Normal workers, so it stays in L2 + // (EVICT_LAST); Shifted consumes it only in the next layer. + ptx::tma_store_3d(&tensor_map_new_residual, smem.residual[io_stage_idx][0], + k_idx, m_idx, 0u, + kIsShifted ? cute::TMA::CacheHintSm100::EVICT_NORMAL + : cute::TMA::CacheHintSm100::EVICT_LAST); + if constexpr (kIsShifted) + ptx::tma_store_2d(&tensor_map_y_bf16, smem.x[io_stage_idx], k_idx, + m_idx, cute::TMA::CacheHintSm100::EVICT_LAST); + + cute::tma_store_arrive(); + // Issue next-task coefficients here to avoid competing with the task's main TMA wave. + if (should_prefetch_next_coeff) + issue_coeff_tma((task_idx + kNumSMs) / kNumSplits * BLOCK_M, next_coeff_stage_idx); + + // Complete the destination writes before releasing this IO stage. + ptx::tma_store_wait<0>(); + smem.empty_io_barriers[io_stage_idx].arrive(); + } + } + + if constexpr (kIsShifted) { + // Arrive at Norm only after both Post workers have committed their row statistics. + smem.full_stats_barrier.wait(task_phase); + const auto x1_sqr_sum_partial = workspace.get_x1_sqr_sum_partial_ptr(task_idx); + float2 row_sqr_sums = {}; + #pragma unroll + for (uint32_t post_worker_idx = 0; post_worker_idx < kNumPostWGs; ++ post_worker_idx) { + row_sqr_sums.x += smem.x1_sqr_sums[post_worker_idx][lane_idx]; + row_sqr_sums.y += smem.x1_sqr_sums[post_worker_idx][lane_idx + 32]; + } + + x1_sqr_sum_partial[lane_idx] = row_sqr_sums.x; + x1_sqr_sum_partial[lane_idx + 32] = row_sqr_sums.y; + __syncwarp(); + if (cute::elect_one_sync()) { + smem.empty_stats_barrier.arrive(); + sched::mega_mhc::Norm::arrive(gmem_split_barriers, m_block_idx); + } + } else if (cute::elect_one_sync()) { + smem.mix_arrival_barrier.arrive(); + } + } + } + } else if (warp_idx < kWorkspaceEpilogueWGIdx * kNumWarpsPerWG) { + // WG1-2: Post workers, one per WG + const auto get_bf16x2_offset = [](const uint32_t& row_idx, const uint32_t& col_idx) { + constexpr uint32_t kNumColsPerRow = kSwizzleMode / sizeof(nv_bfloat162); + constexpr uint32_t kNumColsPerBank = 16 / sizeof(nv_bfloat162); + const uint32_t bank_group_idx = (col_idx / kNumColsPerBank) ^ (row_idx & (kSwizzleMode / 16 - 1)); + return row_idx * kNumColsPerRow + bank_group_idx * kNumColsPerBank + col_idx % kNumColsPerBank; + }; + + // Comb rows use a separate 64B swizzle to distribute each four-float group across banks. + const auto load_comb_coeff = [](const float* row, const uint32_t& row_idx, const uint32_t& coeff_idx) { + const uint32_t bank_group_idx = (coeff_idx / 4) ^ ((row_idx / 2) & 3); + return row[bank_group_idx * 4 + coeff_idx % 4]; + }; + + DG_STATIC_ASSERT(UMMA_K % kNumLanesPerRowGroup == 0, "Invalid Post geometry"); + + const auto run_post_worker = [&](auto post_worker) { + constexpr uint32_t kPostWorkerIdx = decltype(post_worker)::value; + cutlass::arch::warpgroup_reg_alloc(); + const uint32_t col_pair_idx = lane_idx & (kNumLanesPerRowGroup - 1); + const uint32_t first_row_idx = get_first_row_idx(); + const uint32_t row_indices[kNumTokenRowsPerLane] = {first_row_idx, first_row_idx + kNumRowGroupsPerWarp}; + + RingPipeline io_pipeline; + RingPipeline a_pipeline; + RingPipeline coeff_pipeline; + + // Each Post worker maps to one WG; offset them so their A stages interleave. + a_pipeline.advance(kPostWorkerIdx); + + for (uint32_t task_idx = blockIdx.x, task_phase = 0; task_idx < num_tasks; task_idx += kNumSMs, task_phase ^= 1u) { + const uint32_t split_idx = task_idx % kNumSplits; + const uint32_t num_task_k_blocks = get_num_task_k_blocks(split_idx); + + CUTE_TIE_DECL(coeff_pipeline.advance(), coeff_stage_idx, coeff_phase); + float2 hc_norm_sqr_sums[kNumTokenRowsPerLane] = {}; + float2 x1_sqr_sums[kNumTokenRowsPerLane] = {}; + + // Coefficients stay in registers for every hidden block, so release their SMEM stage immediately. + smem.full_coeff_barriers[coeff_stage_idx].wait(coeff_phase); + float2 row_pre_coeff[kNumRoutes]; + float2 row_post_coeff[kNumRoutes]; + #pragma unroll + for (uint32_t route_idx = 0; route_idx < kNumRoutes; ++ route_idx) { + if constexpr (kIsShifted) + row_pre_coeff[route_idx] = {smem.pre_coeff[coeff_stage_idx][row_indices[0]][route_idx], + smem.pre_coeff[coeff_stage_idx][row_indices[1]][route_idx]}; + row_post_coeff[route_idx] = {smem.post_coeff[coeff_stage_idx][row_indices[0]][route_idx], + smem.post_coeff[coeff_stage_idx][row_indices[1]][route_idx]}; + } + + float2 row_comb_coeff[kNumRoutes * kNumRoutes]; + #pragma unroll + for (uint32_t coeff_idx = 0; coeff_idx < kNumRoutes * kNumRoutes; ++ coeff_idx) + row_comb_coeff[coeff_idx] = {load_comb_coeff(smem.comb_coeff[coeff_stage_idx][row_indices[0]], row_indices[0], coeff_idx), + load_comb_coeff(smem.comb_coeff[coeff_stage_idx][row_indices[1]], row_indices[1], coeff_idx)}; + + __syncwarp(); + if (cute::elect_one_sync()) { + cutlass::arch::fence_view_async_shared(); + smem.empty_coeff_barriers[coeff_stage_idx].arrive(); + } + + // Round-robin UMMA atoms across Post workers and retain statistics until the split ends. + #pragma unroll 1 + for (uint32_t k_block_idx = 0; k_block_idx < num_task_k_blocks; ++ k_block_idx) { + // Wait TMA residual and X arrival + CUTE_TIE_DECL(io_pipeline.advance(), io_stage_idx, io_phase); + smem.full_io_barriers[io_stage_idx].wait(io_phase); + + #pragma unroll + for (uint32_t umma_k_idx = kPostWorkerIdx; umma_k_idx < BLOCK_K / UMMA_K; + umma_k_idx += kNumPostWGs) { + const uint32_t hidden_pair_idx = umma_k_idx * (UMMA_K / 2) + + col_pair_idx * (UMMA_K / kNumLanesPerRowGroup / 2); + float2 x[kNumTokenRowsPerLane]; + float2 residual[kNumTokenRowsPerLane][kNumRoutes]; + #pragma unroll + for (uint32_t row_idx = 0; row_idx < kNumTokenRowsPerLane; ++ row_idx) { + x[row_idx] = __bfloat1622float2( + smem.x[io_stage_idx][get_bf16x2_offset(row_indices[row_idx], hidden_pair_idx)]); + #pragma unroll + for (uint32_t input_route_idx = 0; input_route_idx < kNumRoutes; ++ input_route_idx) + residual[row_idx][input_route_idx] = __bfloat1622float2( + smem.residual[io_stage_idx][input_route_idx][get_bf16x2_offset(row_indices[row_idx], hidden_pair_idx)]); + } + + // Wait MMA release of this worker's next A stage + CUTE_TIE_DECL(a_pipeline.advance(kNumPostWGs), a_stage_idx, a_phase); + smem.empty_a_barriers[a_stage_idx].wait(a_phase ^ 1u); + ptx::tcgen05_after_thread_sync(); + + float2 x1[kNumTokenRowsPerLane] = {}; + float2 post_values[kNumRoutes][kNumTokenRowsPerLane]; + #pragma unroll + for (uint32_t row_idx = 0; row_idx < kNumTokenRowsPerLane; ++ row_idx) { + #pragma unroll + for (uint32_t output_route_idx = 0; output_route_idx < kNumRoutes; ++ output_route_idx) { + const float2 post = row_post_coeff[output_route_idx]; + const float post_coeff = row_idx == 0 ? post.x : post.y; + auto& post_value = post_values[output_route_idx][row_idx]; + post_value = __fmul2_rn(x[row_idx], {post_coeff, post_coeff}); + #pragma unroll + for (uint32_t input_route_idx = 0; input_route_idx < kNumRoutes; ++ input_route_idx) { + const uint32_t comb_idx = input_route_idx * kNumRoutes + output_route_idx; + const float2 comb = row_comb_coeff[comb_idx]; + const float coeff = row_idx == 0 ? comb.x : comb.y; + post_value = __ffma2_rn(residual[row_idx][input_route_idx], {coeff, coeff}, post_value); + } + // NOTES: this BF16 round is the shared numerical boundary for all consumers. + const nv_bfloat162 post_value_bf16 = __float22bfloat162_rn(post_value); + post_value = __bfloat1622float2(post_value_bf16); + smem.residual[io_stage_idx][output_route_idx][ + get_bf16x2_offset(row_indices[row_idx], hidden_pair_idx)] = post_value_bf16; + hc_norm_sqr_sums[row_idx] = __ffma2_rn(post_value, post_value, hc_norm_sqr_sums[row_idx]); + if constexpr (kIsShifted) { + const float2 pre = row_pre_coeff[output_route_idx]; + const float coeff = row_idx == 0 ? pre.x : pre.y; + x1[row_idx] = __ffma2_rn(post_value, {coeff, coeff}, x1[row_idx]); + } + } + } + + const auto tmem_values = reinterpret_cast(post_values); + const uint32_t tmem_addr = tmem_base + a_stage_idx * kNumTmemColumnsPerAStage; + cute::SM100_TMEM_STORE_16dp256b4x::copy(tmem_values[0][0], tmem_values[0][1], tmem_values[0][2], tmem_values[0][3], + tmem_values[1][0], tmem_values[1][1], tmem_values[1][2], tmem_values[1][3], + tmem_values[2][0], tmem_values[2][1], tmem_values[2][2], tmem_values[2][3], + tmem_values[3][0], tmem_values[3][1], tmem_values[3][2], tmem_values[3][3], tmem_addr); + + if constexpr (kIsShifted) { + #pragma unroll + for (uint32_t row_idx = 0; row_idx < kNumTokenRowsPerLane; ++ row_idx) { + const nv_bfloat162 x1_bf16 = __float22bfloat162_rn(x1[row_idx]); + // NOTES: X1 is already rounded to BF16; accumulate without unpacking it back to FP32. + ptx::accumulate_square(x1_sqr_sums[row_idx], x1_bf16); + smem.x[io_stage_idx][get_bf16x2_offset(row_indices[row_idx], hidden_pair_idx)] = x1_bf16; + } + } + + // All four warps have issued their route stores before A becomes visible to the MMA warp. + cutlass::arch::fence_view_async_tmem_store(); + ptx::tcgen05_before_thread_sync(); + smem.full_a_barriers[a_stage_idx].arrive(); + } + + // Make generic-proxy writes visible before the output-store warp reads through the TMA async proxy. + __syncwarp(); + if (cute::elect_one_sync()) { + cute::tma_store_fence(); + smem.store_ready_io_barriers[io_stage_idx].arrive(); + } + } + + // Wait only when register-held statistics are committed to shared rows reused by the next task. + smem.empty_stats_barrier.wait(task_phase ^ 1u); + + // Four lanes own both column pairs of a row, independent of token batch size and SM scheduling. + const float reduced_hc_norm_0 = math::warp_reduce_sum(hc_norm_sqr_sums[0].x + hc_norm_sqr_sums[0].y); + const float reduced_hc_norm_1 = math::warp_reduce_sum(hc_norm_sqr_sums[1].x + hc_norm_sqr_sums[1].y); + float reduced_x1_0, reduced_x1_1; + if constexpr (kIsShifted) { + reduced_x1_0 = math::warp_reduce_sum(x1_sqr_sums[0].x + x1_sqr_sums[0].y); + reduced_x1_1 = math::warp_reduce_sum(x1_sqr_sums[1].x + x1_sqr_sums[1].y); + } + + smem.hc_norm_sqr_sums[kPostWorkerIdx][row_indices[0]] = reduced_hc_norm_0; + smem.hc_norm_sqr_sums[kPostWorkerIdx][row_indices[1]] = reduced_hc_norm_1; + if constexpr (kIsShifted) { + smem.x1_sqr_sums[kPostWorkerIdx][row_indices[0]] = reduced_x1_0; + smem.x1_sqr_sums[kPostWorkerIdx][row_indices[1]] = reduced_x1_1; + } + __syncwarp(); + + if (cute::elect_one_sync()) + smem.full_stats_barrier.arrive(); + __syncwarp(); + } + }; + warp_idx < (kFirstPostWGIdx + 1) * kNumWarpsPerWG ? + run_post_worker(cute::Int<0>{}) : run_post_worker(cute::Int<1>{}); + cutlass::arch::warpgroup_reg_dealloc(); + } else if (warp_idx < kMixWGIdx * kNumWarpsPerWG) { + // WG3: workspace-epilogue worker, draining one split accumulator block per grid-stride task + cutlass::arch::warpgroup_reg_dealloc(); + + const uint32_t warp_idx_in_wg = warp_idx % kNumWarpsPerWG; + const uint32_t col_pair_idx = lane_idx & (kNumLanesPerRowGroup - 1); + const uint32_t first_row_idx = get_first_row_idx(); + const uint32_t row_indices[kNumTokenRowsPerLane] = {first_row_idx, first_row_idx + kNumRowGroupsPerWarp}; + + RingPipeline accum_pipeline; + + for (uint32_t task_idx = blockIdx.x, task_phase = 0; task_idx < num_tasks; task_idx += kNumSMs, task_phase ^= 1u) { + const uint32_t m_block_idx = task_idx / kNumSplits; + if (warp_idx_in_wg == 0) + sched::mega_mhc::Mix::wait_init(gmem_split_barriers, m_block_idx); + + const uint32_t token_indices[kNumTokenRowsPerLane] = {m_block_idx * BLOCK_M + row_indices[0], + m_block_idx * BLOCK_M + row_indices[1]}; + const auto gemm_partial = workspace.get_gemm_partial_ptr(task_idx); + + // Wait this split's accumulated HC outputs + CUTE_TIE_DECL(accum_pipeline.advance(), accum_stage_idx, accum_phase); + smem.full_accum_barriers[accum_stage_idx].wait(accum_phase); + ptx::tcgen05_after_thread_sync(); + + // Each group covers eight FP32 outputs per row; the 2x and 1x loads fetch all three groups. + constexpr uint32_t kNumOutputValuesPerLoad = 32 / sizeof(float); + constexpr uint32_t kNumOutputValuesPerLane = kNumOutputValuesPerLoad / kNumLanesPerRowGroup; + DG_STATIC_ASSERT(kNumOutputValuesPerLoad % kNumLanesPerRowGroup == 0, + "Invalid workspace epilogue geometry"); + + DG_STATIC_ASSERT(BLOCK_N == 3 * kNumOutputValuesPerLoad, "Unsupported workspace output width"); + uint32_t values[3 * kNumLanesPerRowGroup]; + const uint32_t tmem_addr = tmem_base + kAccumTmemStartColumn + accum_stage_idx * BLOCK_N; + cute::SM100_TMEM_LOAD_16dp256b2x::copy(tmem_addr, values[0], values[1], values[2], values[3], + values[4], values[5], values[6], values[7]); + cute::SM100_TMEM_LOAD_16dp256b1x::copy(tmem_addr + 2 * kNumOutputValuesPerLoad, + values[8], values[9], values[10], values[11]); + cutlass::arch::fence_view_async_tmem_load(); + + #pragma unroll + for (uint32_t output_idx = 0; output_idx < BLOCK_N; output_idx += kNumOutputValuesPerLoad) { + const uint32_t value_idx = output_idx / kNumOutputValuesPerLane; + float2 row_values[kNumTokenRowsPerLane] = { + {__uint_as_float(values[value_idx]), __uint_as_float(values[value_idx + 1])}, + {__uint_as_float(values[value_idx + 2]), __uint_as_float(values[value_idx + 3])}}; + #pragma unroll + for (uint32_t row_idx = 0; row_idx < kNumTokenRowsPerLane; ++ row_idx) { + if (token_indices[row_idx] >= norm_args.num_tokens) + row_values[row_idx] = {0.0f, 0.0f}; + } + + *reinterpret_cast( + gemm_partial + row_indices[0] * BLOCK_N + output_idx + col_pair_idx * kNumOutputValuesPerLane) = row_values[0]; + *reinterpret_cast( + gemm_partial + row_indices[1] * BLOCK_N + output_idx + col_pair_idx * kNumOutputValuesPerLane) = row_values[1]; + } + + ptx::tcgen05_before_thread_sync(); + if (cute::elect_one_sync()) + smem.empty_accum_barriers[accum_stage_idx].arrive(); + + // Gate this split's Mix arrival on all four epilogue warps and, in Normal, the output-store warp. + smem.full_stats_barrier.wait(task_phase); + + if (lane_idx < kNumRowsPerWarp) { + const uint32_t row_idx = warp_idx_in_wg * kNumRowsPerWarp + lane_idx; + const auto hc_norm_sqr_sum_partial = workspace.get_hc_norm_sqr_sum_partial_ptr(task_idx); + float row_hc_norm_sqr_sum = 0.0f; + #pragma unroll + for (uint32_t post_worker_idx = 0; post_worker_idx < kNumPostWGs; ++ post_worker_idx) + row_hc_norm_sqr_sum += smem.hc_norm_sqr_sums[post_worker_idx][row_idx]; + hc_norm_sqr_sum_partial[row_idx] = row_hc_norm_sqr_sum; + } + + __syncwarp(); + if (cute::elect_one_sync()) { + smem.empty_stats_barrier.arrive(); + smem.mix_arrival_barrier.arrive(); + } + + if (warp_idx_in_wg == 0) { + smem.mix_arrival_barrier.wait(task_phase); + if (cute::elect_one_sync()) + sched::mega_mhc::Mix::arrive(gmem_split_barriers, m_block_idx); + } + } + } else if (warp_idx < kNormWGIdx * kNumWarpsPerWG) { + // WG4: four shifted Mix workers, one per warp + if constexpr (kIsShifted) { + cutlass::arch::warpgroup_reg_dealloc(); + const uint32_t warp_idx_in_mix_wg = warp_idx - kMixWGIdx * kNumWarpsPerWG; + for (uint32_t token_idx = blockIdx.x * kNumWarpsPerWG + warp_idx_in_mix_wg; + token_idx < norm_args.num_tokens; token_idx += kNumSMs * kNumWarpsPerWG) + epilogue::mega_mhc::run_mix_task( + workspace, gmem_split_barriers, mix_args, token_idx, lane_idx); + } + } + + if constexpr (not kIsShifted) { + // Normal repurposes WG4-5 immediately and each producer WG as soon as its static role retires. + if (warp_idx >= kMixWGIdx * kNumWarpsPerWG) + cutlass::arch::warpgroup_reg_dealloc(); + + const uint32_t wg_idx_in_cta = warp_idx / kNumWarpsPerWG; + + // Each Normal worker owns a distinct named barrier. + cutlass::arch::NamedBarrier::sync(128, 1 + wg_idx_in_cta); + cutlass::arch::warpgroup_reg_alloc(); + + epilogue::mega_mhc::run_normal_worker( + smem, workspace, wg_idx_in_cta, warp_idx % kNumWarpsPerWG, lane_idx, + gmem_split_barriers, mix_args, norm_args); + cutlass::arch::warpgroup_reg_dealloc(); + } else { + // WG5 starts immediately; retired Post workers join the same Norm task queue. + const uint32_t wg_idx = warp_idx / kNumWarpsPerWG; + if ((wg_idx >= kFirstPostWGIdx and wg_idx < kWorkspaceEpilogueWGIdx) or wg_idx == kNormWGIdx) { + epilogue::mega_mhc::run_shifted_norm_worker( + smem, workspace, gmem_split_barriers, norm_args, lane_idx); + } + } + + // All roles must retire before the allocating warp releases TMEM. + __syncthreads(); + if (warp_idx == 3) + Allocator().free(tmem_base, kNumTmemCols); +#else + if (blockIdx.x == 0 and threadIdx.x == 0) + DG_DEVICE_ASSERT(false and "This kernel only supports sm_100f"); +#endif +} + +} // namespace deep_gemm +#pragma clang diagnostic pop diff --git a/deep_gemm/include/deep_gemm/impls/sm100_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm100_mqa_logits.cuh index 969e6df8a8..fdef46dc3c 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_mqa_logits.cuh @@ -6,8 +6,11 @@ #include #include +#include +#include #include #include +#include #include #include #include @@ -21,22 +24,6 @@ namespace deep_gemm { -// Ring-buffer counter avoiding `% kNumStages`, which ptxas can lower poorly for TMEM paths -template -struct RingPipeline { - uint32_t stage_idx = 0, phase = 0; - - CUTLASS_DEVICE cute::tuple advance(const uint32_t& step = 1) { - const uint32_t current_stage_idx = stage_idx, current_phase = phase; - stage_idx += step; - if (stage_idx >= kNumStages) { - stage_idx -= kNumStages; - phase ^= 1; - } - return {current_stage_idx, current_phase}; - } -}; - // Convert runtime valid-token count to `cute::Int` so token loops stay compile-time constant template CUTLASS_DEVICE void dispatch_num_block_tokens(const uint32_t& num_block_tokens, Fn&& fn) { @@ -49,9 +36,24 @@ CUTLASS_DEVICE void dispatch_num_block_tokens(const uint32_t& num_block_tokens, } } +// Load heads using power-of-two TMEM chunks. +template +CUTLASS_DEVICE void load_tmem_heads_decomposed(const uint32_t& tmem_col, float* accum) { + constexpr uint32_t kNumRemaining = kNumHeads - kNumLoaded; + if constexpr (kNumRemaining > 0) { + constexpr uint32_t kChunk = kNumRemaining >= 64 ? 64 + : kNumRemaining >= 32 ? 32 + : kNumRemaining >= 16 ? 16 + : kNumRemaining >= 8 ? 8 + : 4; + ptx::tmem_load_32dp32b(tmem_col + kNumLoaded, reinterpret_cast(accum + kNumLoaded)); + load_tmem_heads_decomposed(tmem_col, accum); + } +} + // Shared device core parameterized by dtype and scheduler geometry/addressing template ::infinity(); + if (warp_idx == kSpecWarpStart) { cute::prefetch_tma_descriptor(&tensor_map_q); cute::prefetch_tma_descriptor(&tensor_map_sf_q); @@ -85,27 +88,32 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, static constexpr uint32_t kNumTmemStages = 3; static constexpr uint32_t kNumUTCCPAlignedElems = 128; static constexpr uint32_t UMMA_M = 128; - static constexpr uint32_t UMMA_N = BLOCK_Q * kNumHeads; + static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; + static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); static constexpr uint32_t UMMA_K = kIsFP4 ? 64 : 32; - static constexpr uint32_t kNumSFQ = kIsMXSF ? math::constexpr_align(BLOCK_Q * kNumHeads, kNumUTCCPAlignedElems) : 0; + static constexpr uint32_t kNumSFQ = kIsMXSF ? math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems) : 0; static constexpr uint32_t kNumSFKV = kIsMXSF ? math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems) : 0; - static constexpr uint32_t kRealNumSFQ = BLOCK_Q * kNumHeads; static constexpr uint32_t kNumQKBytesPerToken = kIsFP4 ? (kHeadDim / 2) : kHeadDim; - static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = BLOCK_Q * kNumHeads * kNumQKBytesPerToken; + static constexpr uint32_t SMEM_Q_SIZE_PER_STAGE = BLOCK_QH * kNumQKBytesPerToken; static constexpr uint32_t SMEM_KV_SIZE_PER_STAGE = SPLIT_KV * kNumQKBytesPerToken; - static constexpr uint32_t SMEM_SF_Q_SIZE_PER_STAGE = kIsMXSF ? (kRealNumSFQ * sizeof(int)) : 0; + static constexpr uint32_t SMEM_SF_Q_SIZE_PER_STAGE = kIsMXSF ? (BLOCK_QH * sizeof(int)) : 0; static constexpr uint32_t SMEM_SF_KV_SIZE_PER_STAGE = kIsMXSF ? (kNumSFKV * sizeof(int)) : (SPLIT_KV * sizeof(float)); - static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = BLOCK_Q * kNumHeads * sizeof(reduce_dtype_t); + static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( + kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); + static constexpr uint32_t kNumWeightElementsPerRow = + kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); + static constexpr uint32_t SMEM_WEIGHT_SIZE_PER_STAGE = BLOCK_Q * kNumWeightBytesPerRow; DG_STATIC_ASSERT(kNumSpecializedThreads == 128 and kNumMathThreads % 128 == 0, "Invalid threads"); DG_STATIC_ASSERT(SPLIT_KV == kNumMathWarpGroups * UMMA_M and SPLIT_KV % kNumUTCCPAlignedElems == 0, "Invalid `SPLIT_KV`"); + DG_STATIC_ASSERT(not (kIsCompressedLogits and kCleanLogits), "Compressed logits cannot be cleaned in-kernel"); using SharedStorage = layout::MQALogitsSharedStorage; extern __shared__ __align__(SharedStorage::kSwizzleAlignment) uint8_t smem_buffer[]; auto& smem = *reinterpret_cast(smem_buffer); - constexpr uint32_t kNumAccumTmemCols = BLOCK_Q * kNumHeads * kNumTmemStages; + constexpr uint32_t kNumAccumTmemCols = UMMA_N * kNumTmemStages; constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); constexpr uint32_t kTmemStartColOfSFQ = kNumAccumTmemCols; constexpr uint32_t kTmemStartColOfSFKV = kNumAccumTmemCols + kNumSFQ / 32; @@ -135,8 +143,6 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); __syncthreads(); - uint32_t seq_k_start[BLOCK_Q], seq_k_end[BLOCK_Q]; - RingPipeline q_pipeline; RingPipeline kv_pipeline; RingPipeline tmem_pipeline; @@ -149,7 +155,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, if (warp_idx == kSpecWarpStart) { cutlass::arch::warpgroup_reg_dealloc(); if (cute::elect_one_sync()) { - auto scheduler = make_scheduler(sm_idx, seq_k_start, seq_k_end); + auto scheduler = make_scheduler(sm_idx); // NOTES: split index for paged scheduler, token offset for contiguous-KV scheduler. uint32_t q_block_idx, kv_base, num_kv_splits; while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { @@ -157,12 +163,12 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, smem.empty_q_barriers[q_stage_idx].wait(q_phase ^ 1); const uint32_t q_token_base = scheduler.get_q_tma_token_base(q_block_idx); - tma::copy( + tma::copy( &tensor_map_q, &smem.full_q_barriers[q_stage_idx], smem.smem_q[q_stage_idx], 0, q_token_base * kNumHeads); if constexpr (kIsMXSF) tma::copy(&tensor_map_sf_q, &smem.full_q_barriers[q_stage_idx], smem.smem_sf_q[q_stage_idx], 0, q_token_base); - tma::copy(&tensor_map_weights, &smem.full_q_barriers[q_stage_idx], smem.smem_weights[q_stage_idx], 0, q_token_base); + tma::copy(&tensor_map_weights, &smem.full_q_barriers[q_stage_idx], smem.smem_weights[q_stage_idx], 0, q_token_base); smem.full_q_barriers[q_stage_idx].arrive_and_expect_tx(SMEM_Q_SIZE_PER_STAGE + SMEM_SF_Q_SIZE_PER_STAGE + SMEM_WEIGHT_SIZE_PER_STAGE); } } @@ -170,7 +176,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, } else if (warp_idx == kSpecWarpStart + 1) { cutlass::arch::warpgroup_reg_dealloc(); - auto scheduler = make_scheduler(sm_idx, seq_k_start, seq_k_end); + auto scheduler = make_scheduler(sm_idx); uint32_t cached_kv_page_base = 0; uint32_t cached_kv_page_coord = 0; // NOTES: split index for paged scheduler, token offset for contiguous-KV scheduler. @@ -205,7 +211,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, if (cute::elect_one_sync()) { #pragma unroll for (uint32_t page_idx = 0; page_idx < kNumPagesPerSplit; ++ page_idx) { - tma::copy( + tma::copy( &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], smem.smem_kv[kv_stage_idx] + page_idx * kPageKV * kNumQKBytesPerToken, 0, 0, 1, page_coords[page_idx]); @@ -221,7 +227,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1); const uint32_t kv_tma_offset = scheduler.get_kv_tma_offset(kv_base, kv_split_idx); - tma::copy( + tma::copy( &tensor_map_kv, &smem.full_kv_barriers[kv_stage_idx], smem.smem_kv[kv_stage_idx], 0, kv_tma_offset); tma::copy(&tensor_map_sf_kv, &smem.full_kv_barriers[kv_stage_idx], @@ -247,7 +253,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, auto sf_desc = mma::sm100::make_sf_desc(nullptr); - auto scheduler = make_scheduler(sm_idx, seq_k_start, seq_k_end); + auto scheduler = make_scheduler(sm_idx); // NOTES: split index for paged scheduler, token offset for contiguous-KV scheduler. uint32_t q_block_idx, kv_base, num_kv_splits; while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { @@ -304,7 +310,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, if constexpr (kIsMXSF) { DG_STATIC_ASSERT((not kIsFP4 and kHeadDim == 32) or kHeadDim == 64 or kHeadDim == 128, "Invalid head dim"); - constexpr uint32_t kPackFactor = kIsFP4 ? 2 : 1; + constexpr uint32_t kPackFactor = get_smem_pack_factor(); constexpr uint32_t kQKSwizzleMode = kHeadDim / kPackFactor; using mma_op_t = cute::conditional_t; @@ -335,8 +341,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, } } - asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" - ::"r"(cute::cast_smem_ptr_to_uint(&smem.full_tmem_barriers[tmem_stage_idx]))); + ptx::umma_arrive_no_elect(smem.full_tmem_barriers[tmem_stage_idx]); } } __syncwarp(); @@ -347,40 +352,47 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, } } else if (warp_idx == kSpecWarpStart + 3) { cutlass::arch::warpgroup_reg_dealloc(); + + if constexpr (kCleanLogits) { + const auto cleaner = epilogue::LogitsCleaner(lane_idx); + + // Cleaning always runs grid-stride + auto scheduler = make_scheduler(sm_idx).make_cleaner(sm_idx); + uint32_t q_block_idx, kv_base, num_kv_splits; + while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + const auto coverage_end = cute::min(kv_base + num_kv_splits * SPLIT_KV, logits_stride); + #pragma unroll 1 + for (uint32_t i = 0; i < BLOCK_Q; ++ i) { + const auto row = logits + scheduler.get_logits_row(q_block_idx, i) * static_cast(logits_stride); + cleaner.fill_row(row, 0, kv_base); + cleaner.fill_row(row, coverage_end, logits_stride); + } + } + } } else if (warp_idx < kSpecWarpStart) { cutlass::arch::warpgroup_reg_alloc(); - const auto math_warpgroup_idx = warpgroup_idx; + constexpr bool kLoadSeqBounds = kIsCompressedLogits or kCleanLogits; + uint32_t seq_k_start[kLoadSeqBounds ? BLOCK_Q : 1]; + uint32_t seq_k_end[kLoadSeqBounds ? BLOCK_Q : 1]; + const auto math_warpgroup_idx = warp_idx / 4; const auto math_thread_idx = warp_idx * 32 + lane_idx; DG_STATIC_ASSERT(kNumMathWarpGroups <= kNumTmemStages, "Math warp groups exceed TMEM stages"); tmem_pipeline.advance(math_warpgroup_idx); constexpr bool kIsReduceBF16 = not cute::is_same_v; - DG_STATIC_ASSERT(not kIsReduceBF16 or kNumHeads % 2 == 0, "bf16 weights need even kNumHeads"); - DG_STATIC_ASSERT(kNumHeads == 4 or kNumHeads == 8 or kNumHeads == 16 or kNumHeads == 32 or kNumHeads == 64, - "Unsupported TMEM load size"); + DG_STATIC_ASSERT(kNumHeads % 4 == 0, "Head count must be a multiple of 4"); + DG_STATIC_ASSERT(8 <= UMMA_N and UMMA_N <= 256, "Invalid UMMA_N for MMA"); using weights_reg_dtype_t = cute::conditional_t; constexpr uint32_t kNumWeightsRegPerToken = kIsReduceBF16 ? (kNumHeads / 2) : kNumHeads; weights_reg_dtype_t weights[BLOCK_Q][kNumWeightsRegPerToken]; float accum[kNumHeads]; - auto tmem_load_no_fence = [](auto num_elems_t, const uint32_t& addr, float* load_dst) { - constexpr uint32_t N = decltype(num_elems_t)::value; - using Loader = cute::conditional_t>>>>; - [&](cute::index_sequence) { - Loader::copy(addr, reinterpret_cast(load_dst)[Is]...); - }(cute::make_index_sequence{}); - }; - - auto scheduler = make_scheduler(sm_idx, seq_k_start, seq_k_end); + auto scheduler = make_scheduler(sm_idx); // NOTES: split index for paged scheduler, token offset for contiguous-KV scheduler. uint32_t q_block_idx, kv_base, num_kv_splits; - while (scheduler.next_q_block(q_block_idx, kv_base, num_kv_splits)) { + while (scheduler.template next_q_block( + q_block_idx, kv_base, num_kv_splits, seq_k_start, seq_k_end)) { CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); smem.full_q_barriers[q_stage_idx].wait(q_phase); @@ -389,7 +401,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, #pragma unroll for (uint32_t i = 0; i < kNumValidTokens; ++ i) { - const auto smem_weights_row = smem.smem_weights[q_stage_idx] + i * kNumHeads; + const auto smem_weights_row = smem.smem_weights[q_stage_idx] + i * kNumWeightElementsPerRow; if constexpr (kIsReduceBF16) { // Load two bf16 weights at a time as one packed shared u32 const auto packed_row = reinterpret_cast(smem_weights_row); @@ -424,23 +436,30 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, ptx::tcgen05_after_thread_sync(); // Release KV smem only after UMMA commits TMEM; earlier release races TMA overwrite - if constexpr (not kIsMXSF) + if constexpr (not kIsMXSF) { + cutlass::arch::fence_view_async_shared(); smem.empty_kv_barriers[kv_stage_idx].arrive(); + } #pragma unroll for (uint32_t i = 0; i < kNumValidTokens; ++ i) { uint32_t tmem_addr = tmem_stage_idx * UMMA_N + i * kNumHeads; if constexpr (kNumHeads == 8) { - tmem_load_no_fence(cute::Int{}, tmem_addr, accum); + ptx::tmem_load_32dp32b(tmem_addr, reinterpret_cast(accum)); cutlass::arch::fence_view_async_tmem_load(); } else if constexpr (kNumHeads == 16) { - tmem_load_no_fence(cute::Int{}, tmem_addr, accum); - tmem_load_no_fence(cute::Int{}, tmem_addr + kNumHeads / 2, accum + kNumHeads / 2); + ptx::tmem_load_32dp32b(tmem_addr, reinterpret_cast(accum)); + ptx::tmem_load_32dp32b(tmem_addr + kNumHeads / 2, + reinterpret_cast(accum + kNumHeads / 2)); cutlass::arch::fence_view_async_tmem_load(); - } else { - tmem_load_no_fence(cute::Int{}, tmem_addr, accum); + } else if constexpr (kNumHeads == 32 or kNumHeads == 64) { + ptx::tmem_load_32dp32b(tmem_addr, reinterpret_cast(accum)); + cutlass::arch::fence_view_async_tmem_load(); + ptx::tmem_load_32dp32b(tmem_addr + kNumHeads / 2, + reinterpret_cast(accum + kNumHeads / 2)); cutlass::arch::fence_view_async_tmem_load(); - tmem_load_no_fence(cute::Int{}, tmem_addr + kNumHeads / 2, accum + kNumHeads / 2); + } else { + load_tmem_heads_decomposed(tmem_addr, accum); cutlass::arch::fence_view_async_tmem_load(); } @@ -493,6 +512,10 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, const uint32_t len = seq_k_end[i] - seq_k_start[i]; if (rel_kv < len) logits[q_offset + rel_kv] = result; + } else if constexpr (kCleanLogits) { + const uint32_t rel_kv = kv_offset - seq_k_start[i]; + const uint32_t len = seq_k_end[i] - seq_k_start[i]; + logits[q_offset + kv_offset] = rel_kv < len ? result : kNegInf; } else { logits[q_offset + kv_offset] = result; } @@ -505,6 +528,7 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, else process_q_block(cute::Int{}); + cutlass::arch::fence_view_async_shared(); smem.empty_q_barriers[q_stage_idx].arrive(); } @@ -516,7 +540,8 @@ CUTLASS_DEVICE void sm100_mqa_logits_core_impl(const uint32_t logits_stride, // Unified contiguous-KV entry for FP8 / MXFP4 / MXFP8. template ( - sm_idx, num_q_tokens, num_kv_tokens, cu_seq_len_k_start, cu_seq_len_k_end, seq_k_start, seq_k_end); + const auto make_scheduler = [&](const uint32_t& sm_idx) { + return sched::SM100MQALogitsScheduler( + sm_idx, num_q_tokens, num_kv_tokens, cu_seq_len_k_start, cu_seq_len_k_end, schedule_meta); }; - sm100_mqa_logits_core_impl( @@ -549,7 +576,7 @@ void sm100_mqa_logits(const uint32_t num_q_tokens, const uint32_t num_kv_tokens, } // Unified paged entry for FP8 / MXFP4 / MXFP8. -// Paged scheduler walks (Q-block, chunk) tasks; BLOCK_Q = 128 / kNumHeads +// Paged scheduler walks (Q-block, chunk) tasks. template ( sm_idx, context_lens, schedule_meta, indices, @@ -581,7 +608,8 @@ void sm100_paged_mqa_logits(const uint32_t num_q_tokens_total, }; // Paged uses `kNumSMs = 0`; schedule meta drives the grid stride - sm100_mqa_logits_core_impl( diff --git a/deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh new file mode 100644 index 0000000000..bfe36287b5 --- /dev/null +++ b/deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh @@ -0,0 +1,627 @@ +#pragma once + +#include +#include + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace deep_gemm { + +namespace sparse_mqa_detail { + +using namespace layout::sparse_mqa_logits; + +template +CUTLASS_DEVICE uint32_t get_swizzled_kv_chunk_idx(const uint32_t logical_chunk_idx) { + constexpr uint32_t kMask = kSwizzleMode / 16 - 1; + return (logical_chunk_idx & ~kMask) | ((logical_chunk_idx & kMask) ^ ((logical_chunk_idx >> 3u) & kMask)); +} + +template +struct ContiguousSparseKVAccessor { + static constexpr bool kSupportsContiguousTMA = true; + static constexpr uint32_t kNumQKBytesPerToken = kHeadDim / get_smem_pack_factor(); + using KVBlockRef = uint32_t; + + const uint8_t* kv; + const uint32_t* sf_kv; + const cute::TmaDescriptor* tensor_map_kv; + const cute::TmaDescriptor* tensor_map_sf_kv; + + CUTLASS_HOST_DEVICE ContiguousSparseKVAccessor(const uint8_t* kv, const uint32_t* sf_kv, + const cute::TmaDescriptor* tensor_map_kv, + const cute::TmaDescriptor* tensor_map_sf_kv) + : kv(kv), sf_kv(sf_kv), tensor_map_kv(tensor_map_kv), tensor_map_sf_kv(tensor_map_sf_kv) {} + + CUTLASS_DEVICE void prefetch_tma_descriptors() const { + cute::prefetch_tma_descriptor(tensor_map_kv); + cute::prefetch_tma_descriptor(tensor_map_sf_kv); + } + + template + CUTLASS_DEVICE void copy_contiguous_kv_split(cutlass::arch::ClusterTransactionBarrier& kv_barrier, + cutlass::arch::ClusterTransactionBarrier& sf_barrier, + void* smem_kv, void* smem_sf_kv, + const uint32_t kv_token_start) const { + DG_STATIC_ASSERT(SPLIT_KV % kNumKVTokensPerTMA == 0, "KV split must contain whole TMA tiles"); + // Publish the small SF transfer first so its transpose can overlap the bulk KV transfer. + #pragma unroll + for (uint32_t token_offset = 0; token_offset < SPLIT_KV; token_offset += kNumKVTokensPerTMA) { + tma::copy(tensor_map_sf_kv, &sf_barrier, + static_cast(smem_sf_kv) + token_offset, kv_token_start + token_offset, 0); + } + sf_barrier.arrive_and_expect_tx(SPLIT_KV * sizeof(uint32_t)); + #pragma unroll + for (uint32_t token_offset = 0; token_offset < SPLIT_KV; token_offset += kNumKVTokensPerTMA) { + tma::copy(tensor_map_kv, &kv_barrier, + static_cast(smem_kv) + token_offset * kNumQKBytesPerToken, 0, kv_token_start + token_offset); + } + kv_barrier.arrive_and_expect_tx(SPLIT_KV * kNumQKBytesPerToken); + } + + CUTLASS_DEVICE KVBlockRef resolve_kv_block(const uint32_t physical_kv_block_idx) const { + return physical_kv_block_idx; + } + + CUTLASS_DEVICE const uint8_t* get_kv_block(const KVBlockRef kv_block_ref, + const uint32_t src_lane_idx) const { + const uint32_t kv_token_start = ptx::exchange(kv_block_ref, src_lane_idx); + // Tail reads rely on PyTorch allocator padding; extra logits are ignored + return kv + static_cast(kv_token_start) * kNumQKBytesPerToken; + } + + CUTLASS_DEVICE const uint32_t* get_sf_kv_block(const KVBlockRef kv_block_ref, + const uint32_t src_lane_idx) const { + return sf_kv + ptx::exchange(kv_block_ref, src_lane_idx); + } +}; + +template +struct PagedSparseKVAccessor { + static constexpr bool kSupportsContiguousTMA = false; + static constexpr uint32_t kNumQKBytesPerToken = kHeadDim / get_smem_pack_factor(); + + struct KVBlockRef { + const uint8_t* kv; + const uint32_t* sf_kv; + + CUTLASS_DEVICE KVBlockRef(const uint8_t* kv, const uint32_t* sf_kv): kv(kv), sf_kv(sf_kv) {} + }; + + const uint8_t* fused_kv_cache; + uint32_t kv_page_stride_bytes; + + CUTLASS_HOST_DEVICE PagedSparseKVAccessor(const uint8_t* fused_kv_cache, + const uint32_t kv_page_stride_bytes) + : fused_kv_cache(fused_kv_cache), kv_page_stride_bytes(kv_page_stride_bytes) {} + + CUTLASS_DEVICE KVBlockRef resolve_kv_block(const uint32_t physical_kv_block_idx) const { + constexpr uint32_t kNumKVBlocksPerPage = PAGE_KV / SPARSE_BLOCK_KV; + const uint32_t physical_page_idx = physical_kv_block_idx / kNumKVBlocksPerPage; + const uint32_t kv_block_idx_in_page = physical_kv_block_idx % kNumKVBlocksPerPage; + const auto page = fused_kv_cache + static_cast(physical_page_idx) * kv_page_stride_bytes; + const uint32_t token_idx_in_page = kv_block_idx_in_page * SPARSE_BLOCK_KV; + return KVBlockRef( + page + token_idx_in_page * kNumQKBytesPerToken, + reinterpret_cast(page + PAGE_KV * kNumQKBytesPerToken) + token_idx_in_page + ); + } + + CUTLASS_DEVICE const uint8_t* get_kv_block(const KVBlockRef& kv_block_ref, + const uint32_t src_lane_idx) const { + return ptx::exchange(kv_block_ref.kv, src_lane_idx); + } + + CUTLASS_DEVICE const uint32_t* get_sf_kv_block(const KVBlockRef& kv_block_ref, + const uint32_t src_lane_idx) const { + return ptx::exchange(kv_block_ref.sf_kv, src_lane_idx); + } +}; + +} // namespace sparse_mqa_detail + +template +CUTLASS_DEVICE void sm100_sparse_mqa_logits_core_impl(const uint32_t logits_stride, nv_bfloat16* logits, + const uint8_t* metadata, const cute::TmaDescriptor& tensor_map_q, + const cute::TmaDescriptor& tensor_map_sf_q, const cute::TmaDescriptor& tensor_map_weights, + const KVAccessor& kv_accessor) { + using namespace sparse_mqa_detail; + using namespace layout::sparse_mqa_logits; + + // MMA configs + constexpr bool kIsFP4 = cute::is_same_v; + constexpr uint32_t kPackFactor = get_smem_pack_factor(); + constexpr uint32_t UMMA_M = 128; + constexpr uint32_t UMMA_N = BLOCK_Q * kNumHeads; + constexpr uint32_t UMMA_K = kIsFP4 ? 64 : 32; + constexpr uint32_t SPLIT_KV = kNumMathWarpGroups * UMMA_M; + + // Thread and register configs + constexpr uint32_t kNumThreads = get_num_threads(kNumMathWarpGroups); + constexpr uint32_t kNumWarpGroups = kNumThreads / 128; + constexpr uint32_t kNumMathThreads = kNumMathWarpGroups * 128; + constexpr uint32_t kNumKVCopyThreads = kNumMathWarpGroups * 32; + constexpr uint32_t kNumEntryRegisters = (512 / kNumWarpGroups / 8) * 8; + constexpr uint32_t kNumControlRegisters = 64; + constexpr uint32_t kNumKVCopyRegisters = 64; + constexpr uint32_t kNumMathRegisters = cute::min(120u, ((kNumEntryRegisters * kNumWarpGroups - kNumControlRegisters - + kNumKVCopyRegisters * (kNumKVCopyThreads / 128)) / kNumMathWarpGroups / 8) * 8); + + // Memory configs + using smem_t = SharedStorage; + constexpr uint32_t kNumKVBlocksPerSplit = smem_t::kNumKVBlocksPerSplit; + constexpr uint32_t kNumSFQ = smem_t::kNumSFQ; + constexpr uint32_t kNumSFQCols = kNumSFQ / 32; + constexpr uint32_t kNumSFKVColsPerStage = SPLIT_KV / 32; + constexpr uint32_t kTmemStartColOfSFQ = UMMA_N * kNumTmemStages; + constexpr uint32_t kTmemStartColOfSFKV = kTmemStartColOfSFQ + kNumQStages * kNumSFQCols; + constexpr uint32_t kNumTmemCols = utils::get_num_aligned_tmem_cols(); + + // Template checks + DG_STATIC_ASSERT(BLOCK_Q == 2 and UMMA_N == 64, "Sparse MQA requires BLOCK_Q=2"); + DG_STATIC_ASSERT(kNumMathWarpGroups == 4 or kNumMathWarpGroups == 5, + "Invalid number of math warpgroups"); + DG_STATIC_ASSERT(kNumTmemStages >= kNumMathWarpGroups, "Invalid TMEM stage count"); + DG_STATIC_ASSERT(kNumTmemCols <= 512 and kNumThreads <= 1024, "Sparse MQA resource overflow"); + DG_STATIC_ASSERT(kNumMathRegisters * kNumMathWarpGroups + kNumControlRegisters + + kNumKVCopyRegisters * (kNumKVCopyThreads / 128) <= kNumEntryRegisters * kNumWarpGroups, + "Register reconfiguration exceeds the CTA entry pool"); + + // Thread indices + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + const uint32_t lane_idx = ptx::get_lane_idx(); + const uint32_t warpgroup_idx = warp_idx / 4; + constexpr uint32_t kQAndMetadataWarpIdx = kNumMathWarpGroups * 4; + constexpr uint32_t kSFTransposeWarpIdx = kQAndMetadataWarpIdx + 1; + + // Shared memory + extern __shared__ __align__(1024) uint8_t smem_buffer[]; + auto& smem = *reinterpret_cast(smem_buffer); + const auto kv_splits = reinterpret_cast*>( + metadata + sizeof(MetadataHeader)); + + // Initialization + if (warp_idx == kQAndMetadataWarpIdx and cute::elect_one_sync()) { + cute::prefetch_tma_descriptor(&tensor_map_q); + cute::prefetch_tma_descriptor(&tensor_map_sf_q); + cute::prefetch_tma_descriptor(&tensor_map_weights); + if constexpr (KVAccessor::kSupportsContiguousTMA) + kv_accessor.prefetch_tma_descriptors(); + #pragma unroll + for (uint32_t stage_idx = 0; stage_idx < kNumQStages; ++ stage_idx) { + smem.full_q_barriers[stage_idx].init(1); + smem.full_sf_q_barriers[stage_idx].init(1); + // Released by the producer warp, UMMA warp, and math warpgroups + smem.empty_q_barriers[stage_idx].init(kNumMathThreads + 64); + } + #pragma unroll + for (uint32_t stage_idx = 0; stage_idx < kNumKVStages; ++ stage_idx) { + smem.full_metadata_barriers[stage_idx].init(32); + smem.full_sf_copy_barriers[stage_idx].init(kNumKVCopyThreads); + // The SF transpose warp contributes the final arrival after its UTCCP completes. + smem.full_kv_barriers[stage_idx].init(kNumKVCopyThreads + 1); + smem.empty_kv_barriers[stage_idx].init(kNumMathThreads + 1); + } + #pragma unroll + for (uint32_t stage_idx = 0; stage_idx < kNumTmemStages; ++ stage_idx) { + smem.full_tmem_barriers[stage_idx].init(1); + smem.empty_tmem_barriers[stage_idx].init(128); + } + cutlass::arch::fence_barrier_init(); + } + if (warp_idx == kSFTransposeWarpIdx) + cute::TMEM::Allocator1Sm().allocate(kNumTmemCols, &smem.tmem_ptr_in_smem); + + // Zero padded Q scales for UTCCP + constexpr uint32_t kNumSFQValues = BLOCK_Q * kNumHeads; + constexpr uint32_t kNumSFQPaddingValues = kNumSFQ - kNumSFQValues; + for (uint32_t padding_idx = threadIdx.x; padding_idx < kNumQStages * kNumSFQPaddingValues; + padding_idx += kNumThreads) { + const uint32_t q_stage_idx = padding_idx / kNumSFQPaddingValues; + const uint32_t stage_padding_idx = padding_idx % kNumSFQPaddingValues; + smem.sf_q[q_stage_idx][kNumSFQValues + stage_padding_idx] = 0; + } + __syncthreads(); + + DG_DEVICE_ASSERT(blockDim.x == kNumThreads and gridDim.x == kNumSMs); + DG_TRAP_ONLY_DEVICE_ASSERT(ptx::ld_shared(&smem.tmem_ptr_in_smem) == 0); + cudaGridDependencySynchronize(); + + constexpr uint32_t kUMMAWarpIdx = kQAndMetadataWarpIdx + 2; + constexpr uint32_t kFirstKVCopyWarpIdx = kNumThreads / 32 - kNumMathWarpGroups; + if (warp_idx == kQAndMetadataWarpIdx) { + cutlass::arch::warpgroup_reg_dealloc(); + RingPipeline q_pipeline; + RingPipeline kv_pipeline; + + const auto header = reinterpret_cast(metadata); + DG_DEVICE_ASSERT(header->use_unaligned_ks == kUseUnalignedKs); + const uint32_t num_waves = header->num_waves; + const auto schedule_entries = reinterpret_cast( + kv_splits + header->num_kv_splits); + for (uint32_t wave_idx = 0; wave_idx < num_waves; ++ wave_idx) { + const auto entry = schedule_entries[wave_idx * kNumSMs + blockIdx.x]; + if (entry.kv_split_begin == entry.kv_split_end) + continue; + + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.empty_q_barriers[q_stage_idx].wait(q_phase ^ 1u); + if (cute::elect_one_sync()) { + smem.q_blocks[q_stage_idx] = QBlock( + entry.q_token_base, entry.num_q_tokens, entry.kv_split_end - entry.kv_split_begin); + tma::copy(&tensor_map_q, &smem.full_q_barriers[q_stage_idx], + smem.q[q_stage_idx][0], 0, entry.q_token_base * kNumHeads); + tma::copy(&tensor_map_sf_q, &smem.full_q_barriers[q_stage_idx], + smem.sf_q[q_stage_idx], 0, entry.q_token_base); + tma::copy(&tensor_map_weights, &smem.full_q_barriers[q_stage_idx], + smem.weights[q_stage_idx], 0, entry.q_token_base); + smem.full_q_barriers[q_stage_idx].arrive_and_expect_tx( + BLOCK_Q * kNumHeads * (kHeadDim / kPackFactor + sizeof(uint32_t) + sizeof(nv_bfloat16))); + } + + for (uint32_t kv_split_idx = entry.kv_split_begin; kv_split_idx < entry.kv_split_end; ++ kv_split_idx) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1u); + if (cute::elect_one_sync()) + ptx::cp_async_cg<256>(reinterpret_cast(&kv_splits[kv_split_idx].header), + reinterpret_cast(&smem.kv_split_headers[kv_stage_idx])); + #pragma unroll + for (uint32_t chunk_idx = lane_idx; chunk_idx < kNumKVBlocksPerSplit / 2; chunk_idx += 32) { + ptx::cp_async_cg<256>(reinterpret_cast(kv_splits[kv_split_idx].kv_block_infos) + chunk_idx, + reinterpret_cast(smem.kv_block_infos[kv_stage_idx]) + chunk_idx); + } + cutlass::arch::cpasync_barrier_arrive_noinc(reinterpret_cast(&smem.full_metadata_barriers[kv_stage_idx])); + } + smem.empty_q_barriers[q_stage_idx].arrive(); + } + + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.empty_q_barriers[q_stage_idx].wait(q_phase ^ 1u); + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.empty_kv_barriers[kv_stage_idx].wait(kv_phase ^ 1u); + if (cute::elect_one_sync()) { + smem.q_blocks[q_stage_idx].num_q_tokens = 0; + smem.kv_split_headers[kv_stage_idx].packed_num_kv_blocks = 0; + // Publish the generic sentinel with a regular release arrival + ptx::mbarrier_arrive_count(smem.full_metadata_barriers[kv_stage_idx], 32); + smem.full_q_barriers[q_stage_idx].arrive(); + } + } else if (warp_idx == kSFTransposeWarpIdx) { + cutlass::arch::warpgroup_reg_dealloc(); + const auto transpose_sf = [&](uint32_t* smem_ptr, const uint32_t num_sf) { + for (uint32_t sf_base = 0; sf_base < num_sf; sf_base += kNumUTCCPAlignedElems) { + uint32_t values[4]; + #pragma unroll + for (uint32_t i = 0; i < 4; ++ i) + values[i] = ptx::ld_shared(smem_ptr + sf_base + i * 32 + lane_idx); + __syncwarp(); + ptx::st_shared(smem_ptr + sf_base + lane_idx * 4, values[0], values[1], values[2], values[3]); + } + }; + auto sf_desc = mma::sm100::make_sf_desc(nullptr); + RingPipeline q_pipeline; + RingPipeline kv_pipeline; + while (true) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.full_q_barriers[q_stage_idx].wait(q_phase); + if (ptx::ld_shared(&smem.q_blocks[q_stage_idx].num_q_tokens) == 0) + break; + const uint32_t num_kv_splits = ptx::ld_shared(&smem.q_blocks[q_stage_idx].num_kv_splits); + transpose_sf(smem.sf_q[q_stage_idx], kNumSFQ); + cutlass::arch::fence_view_async_shared(); + if (cute::elect_one_sync()) { + mma::sm100::replace_smem_desc_addr(sf_desc, smem.sf_q[q_stage_idx]); + cute::SM100_UTCCP_4x32dp128bit_1cta::copy(sf_desc, kTmemStartColOfSFQ + q_stage_idx * kNumSFQCols); + ptx::tcgen05_before_thread_sync(); + smem.full_sf_q_barriers[q_stage_idx].arrive(); + } + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; ++ kv_split_idx) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + // The copy warps only signal this after consuming the corresponding metadata. + smem.full_sf_copy_barriers[kv_stage_idx].wait(kv_phase); + transpose_sf(smem.sf_kv[kv_stage_idx], SPLIT_KV); + cutlass::arch::fence_view_async_shared(); + if (cute::elect_one_sync()) { + #pragma unroll + for (uint32_t sf_idx = 0; sf_idx < SPLIT_KV; sf_idx += kNumUTCCPAlignedElems) { + mma::sm100::replace_smem_desc_addr(sf_desc, smem.sf_kv[kv_stage_idx] + sf_idx); + cute::SM100_UTCCP_4x32dp128bit_1cta::copy(sf_desc, + kTmemStartColOfSFKV + kv_stage_idx * kNumSFKVColsPerStage + sf_idx / 32); + } + ptx::tcgen05_before_thread_sync(); + smem.full_kv_barriers[kv_stage_idx].arrive(); + } + } + } + } else if (warp_idx == kUMMAWarpIdx) { + cutlass::arch::warpgroup_reg_dealloc(); + if (cute::elect_one_sync()) { + using mma_op_t = cute::conditional_t; + const auto instr_desc = cute::UMMA::make_instr_desc_block_scaled(); + RingPipeline q_pipeline; + RingPipeline kv_pipeline; + RingPipeline tmem_pipeline; + while (true) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.full_q_barriers[q_stage_idx].wait(q_phase); + if (ptx::ld_shared(&smem.q_blocks[q_stage_idx].num_q_tokens) == 0) + break; + const uint32_t num_kv_splits = ptx::ld_shared(&smem.q_blocks[q_stage_idx].num_kv_splits); + smem.full_sf_q_barriers[q_stage_idx].wait(q_phase); + ptx::tcgen05_after_thread_sync(); + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; ++ kv_split_idx) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.full_kv_barriers[kv_stage_idx].wait(kv_phase); + cutlass::arch::fence_view_async_shared(); + #pragma unroll + for (uint32_t m_idx = 0; m_idx < kNumMathWarpGroups; ++ m_idx) { + CUTE_TIE_DECL(tmem_pipeline.advance(), tmem_stage_idx, tmem_phase); + const uint32_t tmem_addr = tmem_stage_idx * UMMA_N; + smem.empty_tmem_barriers[tmem_stage_idx].wait(tmem_phase ^ 1u); + ptx::tcgen05_after_thread_sync(); + #pragma unroll + for (uint32_t k_idx = 0; k_idx < kHeadDim / UMMA_K; ++ k_idx) { + const uint32_t sf_id = k_idx * kPackFactor; + const auto runtime_instr_desc = mma::sm100::make_runtime_instr_desc_with_sf_id(instr_desc, sf_id, sf_id); + const auto a_desc = mma::sm100::make_umma_desc( + smem.kv[kv_stage_idx][0], m_idx * UMMA_M, k_idx * UMMA_K); + const auto b_desc = mma::sm100::make_umma_desc( + smem.q[q_stage_idx][0], 0, k_idx * UMMA_K); + mma_op_t::fma(a_desc, b_desc, tmem_addr, k_idx, runtime_instr_desc, + kTmemStartColOfSFKV + kv_stage_idx * kNumSFKVColsPerStage + m_idx * 4, + kTmemStartColOfSFQ + q_stage_idx * kNumSFQCols); + } + ptx::umma_arrive_no_elect(smem.full_tmem_barriers[tmem_stage_idx]); + } + ptx::umma_arrive_no_elect(smem.empty_kv_barriers[kv_stage_idx]); + } + ptx::mbarrier_arrive_count(smem.empty_q_barriers[q_stage_idx], 32); + } + } + } else if (warp_idx >= kFirstKVCopyWarpIdx) { + if (warpgroup_idx == kNumMathWarpGroups) + cutlass::arch::warpgroup_reg_dealloc(); + else + cutlass::arch::warpgroup_reg_dealloc(); + constexpr uint32_t kNumKVBlocksPerWarp = UMMA_M / SPARSE_BLOCK_KV; + constexpr uint32_t kNumChunksPerKVBlock = SPARSE_BLOCK_KV * (kHeadDim / kPackFactor) / 16; + constexpr uint32_t kNumSFChunksPerKVBlock = SPARSE_BLOCK_KV * sizeof(uint32_t) / 16; + DG_STATIC_ASSERT(kNumKVBlocksPerWarp % 2 == 0, "Each KV copy warp must process pairs of sparse KV blocks"); + DG_STATIC_ASSERT(kNumKVBlocksPerWarp * kNumSFChunksPerKVBlock == 32, "Each KV copy warp must issue exactly one full-warp SF copy"); + const uint32_t copy_warp_idx = warp_idx - kFirstKVCopyWarpIdx; + RingPipeline kv_pipeline; + while (true) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + smem.full_metadata_barriers[kv_stage_idx].wait(kv_phase); + const uint32_t packed_num_kv_blocks = ptx::ld_shared(&smem.kv_split_headers[kv_stage_idx].packed_num_kv_blocks); + const uint32_t num_kv_blocks = KVSplitHeader::get_num_kv_blocks(packed_num_kv_blocks); + if (num_kv_blocks == 0) + break; + + if constexpr (KVAccessor::kSupportsContiguousTMA) { + if (KVSplitHeader::is_contiguous(packed_num_kv_blocks)) { + // One TMA lane replaces all generic copy warps for a contiguous split + if (copy_warp_idx == 0 and cute::elect_one_sync()) { + const uint32_t kv_token_start = ptx::ld_shared( + &smem.kv_block_infos[kv_stage_idx][0].physical_kv_block_idx); + kv_accessor.template copy_contiguous_kv_split( + smem.full_kv_barriers[kv_stage_idx], smem.full_sf_copy_barriers[kv_stage_idx], + smem.kv[kv_stage_idx][0], smem.sf_kv[kv_stage_idx], + kv_token_start); + } + __syncwarp(); + // Each copy warp releases its share only after consuming the header + const uint32_t num_arrivals = copy_warp_idx == 0 ? 31 : 32; + // Predication keeps the sparse copy path from being duplicated + const bool is_elected = cute::elect_one_sync(); + ptx::mbarrier_arrive_count_pred( + smem.full_kv_barriers[kv_stage_idx], num_arrivals, is_elected); + ptx::mbarrier_arrive_count_pred( + smem.full_sf_copy_barriers[kv_stage_idx], num_arrivals, is_elected); + continue; + } + } + + const uint32_t kv_block_base_in_split = copy_warp_idx * kNumKVBlocksPerWarp; + if (kv_block_base_in_split < num_kv_blocks) { + const auto kv_block_ref = kv_accessor.resolve_kv_block(lane_idx < kNumKVBlocksPerWarp ? + smem.kv_block_infos[kv_stage_idx][kv_block_base_in_split + lane_idx].physical_kv_block_idx : 0); + + // Copy SF first + if constexpr (kUseUnalignedKs) { + // NOTES: unaligned ks keeps KV rows aligned, but may only 4-byte align SF rows + constexpr uint32_t kNumSFBlocksPerIteration = 32 / SPARSE_BLOCK_KV; + #pragma unroll + for (uint32_t sf_kv_block_base = 0; sf_kv_block_base < kNumKVBlocksPerWarp; sf_kv_block_base += kNumSFBlocksPerIteration) { + const uint32_t sf_kv_block_offset = sf_kv_block_base + lane_idx / SPARSE_BLOCK_KV; + const uint32_t token_in_kv_block = lane_idx % SPARSE_BLOCK_KV; + const auto sf_kv_block = kv_accessor.get_sf_kv_block(kv_block_ref, sf_kv_block_offset); + cute::SM80_CP_ASYNC_CACHEALWAYS::copy(sf_kv_block[token_in_kv_block], + smem.sf_kv[kv_stage_idx][(kv_block_base_in_split + sf_kv_block_offset) * SPARSE_BLOCK_KV + token_in_kv_block]); + } + } else { + const uint32_t sf_kv_block_offset = lane_idx / kNumSFChunksPerKVBlock; + const uint32_t chunk_idx = lane_idx % kNumSFChunksPerKVBlock; + const auto sf_kv_block = kv_accessor.get_sf_kv_block(kv_block_ref, sf_kv_block_offset); + ptx::cp_async_cg<64>(reinterpret_cast(sf_kv_block) + chunk_idx, + reinterpret_cast(smem.sf_kv[kv_stage_idx] + + (kv_block_base_in_split + sf_kv_block_offset) * SPARSE_BLOCK_KV) + chunk_idx); + } + cutlass::arch::cpasync_barrier_arrive_noinc(reinterpret_cast(&smem.full_sf_copy_barriers[kv_stage_idx])); + + // Copy sparse KV blocks + #pragma unroll + for (uint32_t kv_block_pair_offset = 0; kv_block_pair_offset < kNumKVBlocksPerWarp; kv_block_pair_offset += 2) { + const uint32_t kv_block_offset = kv_block_pair_offset + lane_idx / 16; + const auto kv_block = kv_accessor.get_kv_block(kv_block_ref, kv_block_offset); + #pragma unroll + for (uint32_t chunk_base = 0; chunk_base < kNumChunksPerKVBlock; chunk_base += 16) { + const uint32_t chunk_idx = chunk_base + lane_idx % 16; + ptx::cp_async_cg<256>(reinterpret_cast(kv_block) + chunk_idx, + reinterpret_cast(smem.kv[kv_stage_idx][0]) + + (kv_block_base_in_split + kv_block_offset) * kNumChunksPerKVBlock + + get_swizzled_kv_chunk_idx(chunk_idx)); + } + } + } else { + cutlass::arch::cpasync_barrier_arrive_noinc(reinterpret_cast(&smem.full_sf_copy_barriers[kv_stage_idx])); + } + cutlass::arch::cpasync_barrier_arrive_noinc(reinterpret_cast(&smem.full_kv_barriers[kv_stage_idx])); + } + } else if (warpgroup_idx < kNumMathWarpGroups) { + cutlass::arch::warpgroup_reg_alloc(); + const uint32_t math_thread_idx = (warp_idx % 4) * 32 + lane_idx; + RingPipeline q_pipeline; + RingPipeline kv_pipeline; + RingPipeline tmem_pipeline; + tmem_pipeline.advance(warpgroup_idx); + + while (true) { + CUTE_TIE_DECL(q_pipeline.advance(), q_stage_idx, q_phase); + smem.full_q_barriers[q_stage_idx].wait(q_phase); + const uint32_t q_token_base = ptx::ld_shared(&smem.q_blocks[q_stage_idx].q_token_base); + const uint32_t num_q_tokens = ptx::ld_shared(&smem.q_blocks[q_stage_idx].num_q_tokens); + const uint32_t num_kv_splits = ptx::ld_shared(&smem.q_blocks[q_stage_idx].num_kv_splits); + if (num_q_tokens == 0) + break; + + nv_bfloat16* output_rows[BLOCK_Q]; + nv_bfloat162 weights[BLOCK_Q][kNumHeads / 2]; + float accum[kNumHeads / 2]; + + #pragma unroll + for (uint32_t q_token_offset = 0; q_token_offset < BLOCK_Q; ++ q_token_offset) { + output_rows[q_token_offset] = logits + (q_token_base + q_token_offset) * static_cast(logits_stride); + if (q_token_offset >= num_q_tokens) + continue; + const auto packed_weights = reinterpret_cast(smem.weights[q_stage_idx] + q_token_offset * kNumHeads); + #pragma unroll + for (uint32_t i = 0; i < kNumHeads / 2; ++ i) + weights[q_token_offset][i] = packed_weights[i]; + } + + for (uint32_t kv_split_idx = 0; kv_split_idx < num_kv_splits; ++ kv_split_idx) { + CUTE_TIE_DECL(kv_pipeline.advance(), kv_stage_idx, kv_phase); + CUTE_TIE_DECL(tmem_pipeline.advance(kNumMathWarpGroups), tmem_stage_idx, tmem_phase); + + // One 128-token tile per math warpgroup + const uint32_t token_in_kv_split = warpgroup_idx * UMMA_M + math_thread_idx; + const uint32_t kv_block_idx_in_split = token_in_kv_split / SPARSE_BLOCK_KV; + const uint32_t token_in_kv_block = token_in_kv_split - kv_block_idx_in_split * SPARSE_BLOCK_KV; + + smem.full_tmem_barriers[tmem_stage_idx].wait(tmem_phase); + ptx::tcgen05_after_thread_sync(); + // Metadata stays in the generic proxy, so no proxy fence is needed before releasing the stage + const uint32_t q0_slot_base = ptx::ld_shared(&smem.kv_split_headers[kv_stage_idx].q0_slot_base); + const uint32_t q1_slot_base = ptx::ld_shared(&smem.kv_split_headers[kv_stage_idx].q1_slot_base); + const uint32_t packed_slot_offsets = ptx::ld_shared(&smem.kv_block_infos[kv_stage_idx][kv_block_idx_in_split].packed_slot_offsets); + smem.empty_kv_barriers[kv_stage_idx].arrive(); + + #pragma unroll + for (uint32_t q_token_offset = 0; q_token_offset < BLOCK_Q; ++ q_token_offset) { + if (q_token_offset >= num_q_tokens) + continue; + const uint32_t tmem_addr = tmem_stage_idx * UMMA_N + q_token_offset * kNumHeads; + auto sum_0 = __floats2bfloat162_rn(0.0f, 0.0f); + auto sum_1 = __floats2bfloat162_rn(0.0f, 0.0f); + #pragma unroll + for (uint32_t head_base = 0; head_base < kNumHeads; head_base += kNumHeads / 2) { + ptx::tmem_load_32dp32b(tmem_addr + head_base, + reinterpret_cast(accum)); + cutlass::arch::fence_view_async_tmem_load(); + if (q_token_offset + 1 == num_q_tokens and head_base == kNumHeads / 2) { + ptx::tcgen05_before_thread_sync(); + smem.empty_tmem_barriers[tmem_stage_idx].arrive(); + } + #pragma unroll + for (uint32_t head_offset = 0; head_offset < kNumHeads / 2; head_offset += 4) { + const auto accum_pair_0 = make_float2(accum[head_offset], accum[head_offset + 1]); + const auto accum_pair_1 = make_float2(accum[head_offset + 2], accum[head_offset + 3]); + sum_0 = __hfma2(ptx::cvt_relu_bf16x2_f32(accum_pair_0), weights[q_token_offset][(head_base + head_offset) / 2], sum_0); + sum_1 = __hfma2(ptx::cvt_relu_bf16x2_f32(accum_pair_1), weights[q_token_offset][(head_base + head_offset + 2) / 2], sum_1); + } + } + const auto sum = __hadd2_rn(sum_0, sum_1); + const nv_bfloat16 reduced = __hadd_rn(sum.x, sum.y); + + // Map split-local slots back to compressed-logits columns + const uint32_t q_slot_offset = (packed_slot_offsets >> (q_token_offset * kNumSparseSlotBits)) & kInvalidSparseSlot; + if (q_slot_offset != kInvalidSparseSlot) { + const uint32_t q_slot_base = q_token_offset == 0 ? q0_slot_base : q1_slot_base; + const uint32_t output_col_idx = (q_slot_base + q_slot_offset) * SPARSE_BLOCK_KV + token_in_kv_block; + output_rows[q_token_offset][output_col_idx] = reduced; + } + } + } + cutlass::arch::fence_view_async_shared(); + smem.empty_q_barriers[q_stage_idx].arrive(); + } + cutlass::arch::NamedBarrier(kNumMathThreads, 0).sync(); + if (warp_idx == 0) + cute::TMEM::Allocator1Sm().free(0, kNumTmemCols); + } else { + cutlass::arch::warpgroup_reg_dealloc(); + } +} + +template +CUTLASS_GLOBAL __launch_bounds__(layout::sparse_mqa_logits::get_num_threads(kNumMathWarpGroups), 1) +void sm100_sparse_mqa_logits(const uint32_t logits_stride, nv_bfloat16* logits, + const uint8_t* kv, const uint32_t* sf_kv, + const uint8_t* metadata, const __grid_constant__ cute::TmaDescriptor tensor_map_q, + const __grid_constant__ cute::TmaDescriptor tensor_map_sf_q, + const __grid_constant__ cute::TmaDescriptor tensor_map_weights, + const __grid_constant__ cute::TmaDescriptor tensor_map_kv, + const __grid_constant__ cute::TmaDescriptor tensor_map_sf_kv) { + // Keep the CUTLASS dtype out of the kernel template signature to avoid ptxas register spills + using qk_dtype_t = cute::conditional_t; + const auto kv_accessor = sparse_mqa_detail::ContiguousSparseKVAccessor( + kv, sf_kv, &tensor_map_kv, &tensor_map_sf_kv); + sm100_sparse_mqa_logits_core_impl( + logits_stride, logits, metadata, tensor_map_q, tensor_map_sf_q, tensor_map_weights, kv_accessor); +} + +template +CUTLASS_GLOBAL __launch_bounds__(layout::sparse_mqa_logits::get_num_threads(kNumMathWarpGroups), 1) +void sm100_paged_sparse_mqa_logits(const uint32_t logits_stride, const uint32_t kv_page_stride_bytes, + nv_bfloat16* logits, + const uint8_t* fused_kv_cache, const uint8_t* metadata, + const __grid_constant__ cute::TmaDescriptor tensor_map_q, + const __grid_constant__ cute::TmaDescriptor tensor_map_sf_q, + const __grid_constant__ cute::TmaDescriptor tensor_map_weights) { + // Keep the CUTLASS dtype out of the kernel template signature to avoid ptxas register spills + using qk_dtype_t = cute::conditional_t; + DG_STATIC_ASSERT(PAGE_KV % SPARSE_BLOCK_KV == 0, "Sparse KV blocks must not cross pages"); + DG_DEVICE_ASSERT(kv_page_stride_bytes % 512 == 0); + const auto kv_accessor = sparse_mqa_detail::PagedSparseKVAccessor( + fused_kv_cache, kv_page_stride_bytes); + sm100_sparse_mqa_logits_core_impl( + logits_stride, logits, metadata, tensor_map_q, tensor_map_sf_q, tensor_map_weights, kv_accessor); +} + +} // namespace deep_gemm diff --git a/deep_gemm/include/deep_gemm/impls/sm100_tf32_hc_prenorm_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm100_tf32_hc_prenorm_gemm.cuh index aaf7fd9aea..76a06b9000 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_tf32_hc_prenorm_gemm.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_tf32_hc_prenorm_gemm.cuh @@ -324,6 +324,7 @@ sm100_tf32_hc_prenorm_gemm_impl(const uint32_t shape_m, cast_stage_idx * BLOCK_K + i * 8); } cutlass::arch::fence_view_async_tmem_store(); + cutlass::arch::fence_view_async_shared(); // Arrive for issuing MMAs ptx::tcgen05_before_thread_sync(); diff --git a/deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh index c103e543ce..ba165ea2f4 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh @@ -156,6 +156,9 @@ sm90_bf16_gemm_impl(int* grouped_layout, // Persistently schedule over blocks while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + if (is_k_grouped_contiguous(kGemmType) and scheduler.current_shape_k == 0) + continue; + // Assign TMA multicast number into A and B // NOTES: there may be additional odd rows/columns or cases where multicast is not possible. const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx); @@ -220,6 +223,9 @@ sm90_bf16_gemm_impl(int* grouped_layout, const uint32_t b_desc_lo = __shfl_sync(0xffffffff, b_desc.reg32_[0], 0); while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + if (is_k_grouped_contiguous(kGemmType) and scheduler.current_shape_k == 0) + continue; + constexpr uint32_t WAVE_BLOCK_M = BLOCK_M <= WGMMA::M ? BLOCK_M : WGMMA::M * 2; DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes"); float accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0}; diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh index b5cfbe5c90..a8453ccbca 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d1d.cuh @@ -177,6 +177,9 @@ sm90_fp8_gemm_1d1d_impl(__nv_fp8_e4m3* gmem_a_ptr, __nv_fp8_e4m3* gmem_b_ptr, // Persistently schedule over blocks while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + if (kGemmType == GemmType::KGroupedContiguous and scheduler.current_shape_k == 0) + continue; + // Assign TMA multicast number into A and B // NOTES: there may be additional odd rows/columns or cases where multicast is not possible. const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx); @@ -192,7 +195,7 @@ sm90_fp8_gemm_1d1d_impl(__nv_fp8_e4m3* gmem_a_ptr, __nv_fp8_e4m3* gmem_b_ptr, last_group_idx = scheduler.current_group_idx; // Directly update current tensor map - const uint64_t current_k_offset = scheduler.current_k_cumsum; + const uint64_t current_k_offset = scheduler.current_k_start; ptx::tensor_map_replace_global_addr_in_smem(smem_tensor_map_a, gmem_a_ptr + current_k_offset * shape_m); ptx::tensor_map_replace_global_addr_in_smem(smem_tensor_map_b, gmem_b_ptr + current_k_offset * shape_n); ptx::tensor_map_replace_global_inner_dim_stride_in_smem(smem_tensor_map_a, scheduler.current_shape_k, scheduler.current_shape_k); @@ -223,7 +226,7 @@ sm90_fp8_gemm_1d1d_impl(__nv_fp8_e4m3* gmem_a_ptr, __nv_fp8_e4m3* gmem_b_ptr, // Issue TMA auto& full_barrier = *full_barriers[stage_idx]; const uint32_t k_idx = k_block_idx * BLOCK_K; - const uint32_t sf_k_idx = scheduler.current_sf_k_cumsum + k_block_idx; + const uint32_t sf_k_idx = scheduler.current_k_start / BLOCK_K + k_block_idx; const auto tensor_map_a_ptr = (kGemmType == GemmType::KGroupedContiguous ? gmem_tensor_map_a : &tensor_map_a_base); const auto tensor_map_b_ptr = (kGemmType == GemmType::KGroupedContiguous ? gmem_tensor_map_b : &tensor_map_b_base); tma::copy(&tensor_map_sfa, &full_barrier, smem_sfa[stage_idx], m_idx, sf_k_idx, num_tma_multicast_a); @@ -254,6 +257,9 @@ sm90_fp8_gemm_1d1d_impl(__nv_fp8_e4m3* gmem_a_ptr, __nv_fp8_e4m3* gmem_b_ptr, // Persistently schedule over blocks while (scheduler.get_next_block(m_block_idx, n_block_idx)) { + if (kGemmType == GemmType::KGroupedContiguous and scheduler.current_shape_k == 0) + continue; + // Accumulation for WGMMA or CUDA promotion DG_STATIC_ASSERT(BLOCK_M == WGMMA::M * (BLOCK_M <= 64 ? 1 : 2), "Invalid block sizes"); const uint32_t current_shape_k = (kGemmType == GemmType::KGroupedContiguous ? scheduler.current_shape_k : shape_k); 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 5a61743bbf..08e7cba66b 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 @@ -45,7 +45,7 @@ template + typename epilogue_op_t> CUTLASS_GLOBAL __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, uint32_t shape_m, uint32_t shape_n, uint32_t shape_k, @@ -430,7 +430,7 @@ sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout, if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) { auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N; auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M; - auto n_idx = epilogue_type_t::apply_index_n(n_block_idx * BLOCK_N + in_block_n_offset); + auto n_idx = epilogue_op_t::apply_index_n(n_block_idx * BLOCK_N + in_block_n_offset); auto m_idx = scheduler.get_global_idx(shape_m, BLOCK_M, m_block_idx); if constexpr (kGemmType == GemmType::Batched) { cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_ptr, diff --git a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh index 225af44168..0c12cc69b7 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -20,7 +21,7 @@ namespace deep_gemm { template arrive(); + if constexpr (kCleanLogits) { + static constexpr uint32_t kNumWarpsPerRow = kNumMathThreads / 32 / BLOCK_Q; + DG_STATIC_ASSERT((kNumMathThreads / 32) % BLOCK_Q == 0, "Invalid warp assignment"); + cutlass::arch::NamedBarrier::sync(kNumMathThreads, 1); + const auto& row_idx = warp_idx / kNumWarpsPerRow; + const auto& q_idx = min(block_q_idx * BLOCK_Q + row_idx, seq_len - 1); + const auto& fill_ks = min(cu_seq_len_k_start[q_idx], stride_logits); + const auto& fill_ke = cu_seq_len_k_end[q_idx]; + + const auto cleaner = epilogue::LogitsCleaner( + (warp_idx % kNumWarpsPerRow) * 32 + lane_idx); + const auto row = logits + (block_q_idx * BLOCK_Q + row_idx) * static_cast(stride_logits); + cleaner.fill_row(row, 0, fill_ks); + cleaner.fill_row(row, min(fill_ke, stride_logits), stride_logits); + } + // Jump to the next block CUTE_TIE(get_next_block_q_idx(), block_q_idx, q_iter_idx); } diff --git a/deep_gemm/include/deep_gemm/impls/smxx_clean_logits.cuh b/deep_gemm/include/deep_gemm/impls/smxx_clean_logits.cuh deleted file mode 100644 index 2f66b980c5..0000000000 --- a/deep_gemm/include/deep_gemm/impls/smxx_clean_logits.cuh +++ /dev/null @@ -1,74 +0,0 @@ -#pragma once - -#include -#include - -#include -#include - -namespace deep_gemm { - -template -CUTLASS_GLOBAL __launch_bounds__(kNumWarps * 32, 1) -void smxx_clean_logits(const uint32_t seq_len, const uint32_t seq_len_kv, const uint64_t stride_logits, - const uint32_t* cu_seq_len_k_start, const uint32_t* cu_seq_len_k_end, logits_dtype_t* logits) { - const uint32_t num_sms = gridDim.x; - const uint32_t sm_idx = blockIdx.x; - const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0); - - constexpr uint32_t kAlignment = 16 / sizeof(logits_dtype_t); - const logits_dtype_t neg_inf = -cute::numeric_limits::infinity(); - - // Allocate filled `-inf` shared memory - extern __shared__ __align__(1024) logits_dtype_t smem_buffer[]; - #pragma unroll - for (uint32_t i = threadIdx.x; i < BLOCK_KV; i += kNumWarps * 32) - smem_buffer[i] = neg_inf; - cute::tma_store_fence(); - __syncthreads(); - - // Assign sequence to each warp - const auto assign_task = [&](const uint32_t& num, const uint32_t& idx, - const uint32_t& start, const uint32_t& total) -> cute::tuple { - const auto per = total / num, rem = total % num; - return {start + idx * per + cute::min(idx, rem), per + (idx < rem)}; - }; - CUTE_TIE_DECL(assign_task(num_sms, sm_idx, 0, seq_len), sm_seq_start, sm_seq_len); - CUTE_TIE_DECL(assign_task(kNumWarps, warp_idx, sm_seq_start, sm_seq_len), warp_seq_start, warp_seq_len); - - // Wait for primary kernel completion - cudaGridDependencySynchronize(); - - if (cute::elect_one_sync()) { - for (uint32_t i = warp_seq_start; i < warp_seq_start + warp_seq_len; ++ i) { - const auto ks = cu_seq_len_k_start == nullptr ? 0 : cu_seq_len_k_start[i / kNextN]; - const auto ke = cu_seq_len_k_end[i / kNextN] - kNextN + i % kNextN + 1; - const auto aligned_ks = ks / kAlignment * kAlignment, aligned_ke = (ke + kAlignment - 1) / kAlignment * kAlignment; - - for (uint32_t left = 0; left < seq_len_kv; left += BLOCK_KV) { - const auto right = cute::min(left + BLOCK_KV, static_cast(stride_logits)); - if (right <= ks or ke <= left) { - cute::SM90_BULK_COPY_S2G::copy(smem_buffer, logits + i * stride_logits + left, (right - left) * sizeof(logits_dtype_t)); - } else { - if (left < aligned_ks) - cute::SM90_BULK_COPY_S2G::copy(smem_buffer, logits + i * stride_logits + left, (aligned_ks - left) * sizeof(logits_dtype_t)); - if (aligned_ke < right) - cute::SM90_BULK_COPY_S2G::copy(smem_buffer, logits + i * stride_logits + aligned_ke, (right - aligned_ke) * sizeof(logits_dtype_t)); - } - } - } - } - __syncwarp(); - - for (uint32_t i = warp_seq_start; i < warp_seq_start + warp_seq_len; ++ i) { - const auto ks = cu_seq_len_k_start == nullptr ? 0 : cu_seq_len_k_start[i / kNextN]; - const auto ke = cu_seq_len_k_end[i / kNextN] - kNextN + i % kNextN + 1; - const auto aligned_ks = ks / kAlignment * kAlignment, aligned_ke = (ke + kAlignment - 1) / kAlignment * kAlignment; - for (uint32_t j = aligned_ks; j < ks; ++ j) - logits[i * stride_logits + j] = neg_inf; - for (uint32_t j = ke; j < aligned_ke; ++ j) - logits[i * stride_logits + j] = neg_inf; - } -} - -} diff --git a/deep_gemm/include/deep_gemm/impls/smxx_layout.cuh b/deep_gemm/include/deep_gemm/impls/smxx_layout.cuh index bf9495e9bf..57b11f34ce 100644 --- a/deep_gemm/include/deep_gemm/impls/smxx_layout.cuh +++ b/deep_gemm/include/deep_gemm/impls/smxx_layout.cuh @@ -7,51 +7,73 @@ namespace deep_gemm { -template -CUTLASS_GLOBAL void transpose_fp32(const float* sf, float* out, const uint32_t mn) { +CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 4) +void transpose_fp32(const float* sf, float* out, const uint32_t mn) { typedef typename utils::Vectorized::vec_t in_vec_t; constexpr static uint32_t kNumElemsPerVec = sizeof(in_vec_t) / sizeof(float); - constexpr static uint32_t SF_VEC_K = SF_K / kNumElemsPerVec; + constexpr static uint32_t kSmemSFK = SF_K < BLOCK_SF_K ? PADDED_SF_K : BLOCK_SF_K + 1; + DG_STATIC_ASSERT(BLOCK_SF_K % kNumElemsPerVec == 0, "Invalid SF block size"); // Shapes and strides extern __shared__ float smem_buffer[]; constexpr auto kNumTMAAlignedElems = static_cast(16 / sizeof(float)); - const auto in_block_mn = min(BLOCK_MN, mn - blockIdx.x * BLOCK_MN); + const auto num_mn_blocks = math::ceil_div(mn, BLOCK_MN); + const auto block_mn_idx = SF_K < BLOCK_SF_K ? blockIdx.x : blockIdx.x % num_mn_blocks; + const auto block_sf_k_idx = SF_K < BLOCK_SF_K ? 0u : blockIdx.x / num_mn_blocks; + const auto block_sf_k_start = block_sf_k_idx * BLOCK_SF_K; + const auto in_block_mn = min(BLOCK_MN, mn - block_mn_idx * BLOCK_MN); + const auto in_block_sf_k = SF_K < BLOCK_SF_K ? SF_K : min(BLOCK_SF_K, SF_K - block_sf_k_start); const auto tma_aligned_mn = math::align(mn, kNumTMAAlignedElems); // Shift into the block sf = sf + static_cast(blockIdx.y) * mn * SF_K; out = out + static_cast(blockIdx.y) * tma_aligned_mn * SF_K; - const auto& local_sf = reinterpret_cast(sf + static_cast(blockIdx.x) * (BLOCK_MN * SF_K)); + const auto local_sf = sf + static_cast(block_mn_idx) * BLOCK_MN * SF_K + block_sf_k_start; // Wait for primary kernel completion cudaGridDependencySynchronize(); // Load - for (uint32_t i = threadIdx.x; i < in_block_mn * SF_VEC_K; i += kNumThreads) { - auto in_vec = local_sf[i]; - const auto& in_values = reinterpret_cast(&in_vec); - - const auto& row = i / SF_VEC_K, col = (i % SF_VEC_K) * kNumElemsPerVec; - #pragma unroll - for (uint32_t j = 0; j < kNumElemsPerVec; ++ j) - smem_buffer[row * PADDED_SF_K + col + j] = in_values[j]; + if constexpr (SF_K < BLOCK_SF_K) { + constexpr auto kNumSFKVecs = SF_K / kNumElemsPerVec; + for (uint32_t i = threadIdx.x; i < in_block_mn * kNumSFKVecs; i += kNumThreads) { + auto in_vec = reinterpret_cast(local_sf)[i]; + const auto& in_values = reinterpret_cast(&in_vec); + const auto row = i / kNumSFKVecs, col = (i % kNumSFKVecs) * kNumElemsPerVec; + #pragma unroll + for (uint32_t j = 0; j < kNumElemsPerVec; ++ j) + smem_buffer[row * kSmemSFK + col + j] = in_values[j]; + } + } else { + const auto warp_idx = threadIdx.x / 32; + const auto lane_idx = ptx::get_lane_idx(); + for (uint32_t row = warp_idx; row < in_block_mn; row += kNumThreads / 32) { + for (uint32_t col = lane_idx * kNumElemsPerVec; col < in_block_sf_k; col += 32 * kNumElemsPerVec) { + auto in_vec = *reinterpret_cast(local_sf + row * SF_K + col); + const auto& in_values = reinterpret_cast(&in_vec); + #pragma unroll + for (uint32_t j = 0; j < kNumElemsPerVec; ++ j) + smem_buffer[row * kSmemSFK + col + j] = in_values[j]; + } + } } __syncthreads(); // Store #pragma unroll - for (uint32_t i = threadIdx.x; i < in_block_mn * SF_K; i += kNumThreads) { + for (uint32_t i = threadIdx.x; i < in_block_mn * in_block_sf_k; i += kNumThreads) { const auto& sf_k_idx = i / in_block_mn, mn_idx = i % in_block_mn; - const auto& global_mn_idx = blockIdx.x * BLOCK_MN + mn_idx; - out[sf_k_idx * tma_aligned_mn + global_mn_idx] = ptx::ld_shared(smem_buffer + mn_idx * PADDED_SF_K + sf_k_idx); + const auto& global_mn_idx = block_mn_idx * BLOCK_MN + mn_idx; + const auto& global_sf_k_idx = block_sf_k_start + sf_k_idx; + out[global_sf_k_idx * tma_aligned_mn + global_mn_idx] = ptx::ld_shared(smem_buffer + mn_idx * kSmemSFK + sf_k_idx); } } // NOTES: the two kernels below always pack the K dimension -template CUTLASS_GLOBAL void transpose_and_pack_fp32_into_ue8m0(float* sf, uint32_t* out, const uint32_t mn, const uint32_t* grouped_layout, const uint32_t m_alignment) { @@ -60,7 +82,17 @@ CUTLASS_GLOBAL void transpose_and_pack_fp32_into_ue8m0(float* sf, uint32_t* out, // Shapes and strides constexpr auto kNumPackedSFK = math::constexpr_ceil_div(SF_K, 4u); constexpr auto kNumTMAAlignedElems = static_cast(16 / sizeof(int)); - const auto in_block_mn = min(BLOCK_MN, mn - blockIdx.x * BLOCK_MN); + typedef typename utils::Vectorized::vec_t in_vec_t; + constexpr auto kNumElemsPerVec = sizeof(in_vec_t) / sizeof(float); + constexpr auto kSmemSFK = SF_K < BLOCK_SF_K ? SF_K : BLOCK_SF_K + 1; + DG_STATIC_ASSERT(BLOCK_SF_K % 4 == 0 and BLOCK_SF_K % kNumElemsPerVec == 0, "Invalid SF block size"); + const auto num_mn_blocks = math::ceil_div(mn, BLOCK_MN); + const auto block_mn_idx = SF_K < BLOCK_SF_K ? blockIdx.x : blockIdx.x % num_mn_blocks; + const auto block_sf_k_idx = SF_K < BLOCK_SF_K ? 0u : blockIdx.x / num_mn_blocks; + const auto block_sf_k_start = block_sf_k_idx * BLOCK_SF_K; + const auto in_block_mn = min(BLOCK_MN, mn - block_mn_idx * BLOCK_MN); + const auto in_block_sf_k = SF_K < BLOCK_SF_K ? SF_K : min(BLOCK_SF_K, SF_K - block_sf_k_start); + const auto in_block_packed_sf_k = math::ceil_div(in_block_sf_k, 4u); const auto tma_aligned_mn = math::align(mn, kNumTMAAlignedElems); // Wait for primary kernel completion before reading SFs or PSUM layout @@ -71,7 +103,7 @@ CUTLASS_GLOBAL void transpose_and_pack_fp32_into_ue8m0(float* sf, uint32_t* out, const auto group_mn_end = smem_buffer + kNumPsumGroups; const auto sf_smem_buffer = smem_buffer + kNumPsumLayoutElems; - // Precompute PSUM valid MN ranges into smem before the SF tile. + // Precompute PSUM valid MN ranges into smem before the SF block. if constexpr (kUsePsumLayout) { for (uint32_t g = threadIdx.x; g < kNumPsumGroups; g += kNumThreads) { const auto end_g = grouped_layout[g]; @@ -99,25 +131,39 @@ CUTLASS_GLOBAL void transpose_and_pack_fp32_into_ue8m0(float* sf, uint32_t* out, // Load FP32 SFs DG_STATIC_ASSERT(BLOCK_MN % 4 == 0, "Invalid block size"); - const auto local_sf = reinterpret_cast(sf + static_cast(blockIdx.x) * (BLOCK_MN * SF_K)); - const auto num_values = in_block_mn * SF_K; - const auto num_uint4 = num_values / 4; - #pragma unroll - for (uint32_t i = threadIdx.x; i < num_uint4; i += kNumThreads) { - const auto& [x, y, z, w] = reinterpret_cast(local_sf)[i]; - ptx::st_shared(reinterpret_cast(sf_smem_buffer) + i, x, y, z, w); + const auto local_sf = sf + static_cast(block_mn_idx) * BLOCK_MN * SF_K + block_sf_k_start; + if constexpr (SF_K < BLOCK_SF_K) { + const auto local_sf_u32 = reinterpret_cast(local_sf); + const auto num_values = in_block_mn * SF_K; + const auto num_uint4 = num_values / 4; + #pragma unroll + for (uint32_t i = threadIdx.x; i < num_uint4; i += kNumThreads) { + const auto& [x, y, z, w] = reinterpret_cast(local_sf_u32)[i]; + ptx::st_shared(reinterpret_cast(sf_smem_buffer) + i, x, y, z, w); + } + if (const auto unaligned_idx = num_uint4 * 4 + threadIdx.x; unaligned_idx < num_values) + ptx::st_shared(sf_smem_buffer + unaligned_idx, local_sf_u32[unaligned_idx]); + } else { + const auto warp_idx = threadIdx.x / 32; + const auto lane_idx = ptx::get_lane_idx(); + for (uint32_t row = warp_idx; row < in_block_mn; row += kNumThreads / 32) { + for (uint32_t col = lane_idx * kNumElemsPerVec; col < in_block_sf_k; col += 32 * kNumElemsPerVec) { + const auto in_vec = *reinterpret_cast(local_sf + row * SF_K + col); + const auto in_values = reinterpret_cast(&in_vec); + #pragma unroll + for (uint32_t j = 0; j < kNumElemsPerVec; ++ j) + ptx::st_shared(sf_smem_buffer + row * kSmemSFK + col + j, in_values[j]); + } + } } - - // Fill unaligned values as well - if (const auto unaligned_idx = num_uint4 * 4 + threadIdx.x; unaligned_idx < num_values) - ptx::st_shared(sf_smem_buffer + unaligned_idx, local_sf[unaligned_idx]); __syncthreads(); // Pack into UE8M0 and store #pragma unroll - for (uint32_t i = threadIdx.x; i < (kNumPackedSFK * BLOCK_MN); i += kNumThreads) { + for (uint32_t i = threadIdx.x; i < in_block_packed_sf_k * BLOCK_MN; i += kNumThreads) { const auto sf_k_pack_idx = i / BLOCK_MN, mn_idx = i % BLOCK_MN; - const auto global_mn_idx = blockIdx.x * BLOCK_MN + mn_idx; + const auto global_mn_idx = block_mn_idx * BLOCK_MN + mn_idx; + const auto global_sf_k_pack_idx = block_sf_k_start / 4 + sf_k_pack_idx; const auto in_bounds_mn = global_mn_idx < mn; const auto valid_mn = is_valid_mn(global_mn_idx); @@ -126,9 +172,10 @@ CUTLASS_GLOBAL void transpose_and_pack_fp32_into_ue8m0(float* sf, uint32_t* out, #pragma unroll for (uint32_t j = 0; j < 4; ++ j) { const auto sf_k_idx = sf_k_pack_idx * 4 + j; - values[j] = valid_mn and sf_k_idx < SF_K ? ptx::ld_shared(sf_smem_buffer + mn_idx * SF_K + sf_k_idx) : 0; + const auto global_sf_k_idx = block_sf_k_start + sf_k_idx; + values[j] = valid_mn and global_sf_k_idx < SF_K ? ptx::ld_shared(sf_smem_buffer + mn_idx * kSmemSFK + sf_k_idx) : 0; // FP32 SFs must have a zero sign and mantissa (only the exponent is packed) - DG_DEVICE_ASSERT((values[j] & 0x807fffffu) == 0); + DG_TRAP_ONLY_DEVICE_ASSERT((values[j] & 0x807fffffu) == 0); } // Pack and store @@ -139,7 +186,7 @@ CUTLASS_GLOBAL void transpose_and_pack_fp32_into_ue8m0(float* sf, uint32_t* out, packed |= (values[3] << 1u); // Write safe finite scale codes for PSUM gap rows; UE8M0 0xff is NaN. if (in_bounds_mn) - out[sf_k_pack_idx * tma_aligned_mn + global_mn_idx] = packed; + out[global_sf_k_pack_idx * tma_aligned_mn + global_mn_idx] = packed; } } @@ -150,7 +197,7 @@ CUTLASS_GLOBAL void pack_fp32_into_ue8m0(float* sf, uint32_t* out, uint32_t* gro const uint32_t gran_k, const uint32_t k_alignment) { // Always packing the K dimension // NOTES: should also assert `mn % 4 == 0` at launch - // psum layout may have input gaps, but packed output remains per-group compact. + // Each group starts at a new packed row; unused tail slots are zero-filled DG_STATIC_ASSERT(kTransposed, "Currently only support transposed SFs (MN-major)"); DG_STATIC_ASSERT(BLOCK_MN % 4 == 0, "Invalid block sizes"); DG_STATIC_ASSERT(BLOCK_PACKED_SF_K == kNumThreads / 32, "Invalid block sizes"); @@ -178,7 +225,6 @@ CUTLASS_GLOBAL void pack_fp32_into_ue8m0(float* sf, uint32_t* out, uint32_t* gro uint32_t num_padding_sf_rows = 0; uint32_t group_sf_row_start = 0, group_sf_row_end = sf_k; if constexpr (kNumGroups > 1 or kUsePsumLayout) { - // Layout entries are group K sizes, or psum end offsets in psum mode DG_STATIC_ASSERT(kNumGroups <= 128, "Too many groups"); uint32_t layout_cache[4]; #pragma unroll @@ -201,7 +247,8 @@ CUTLASS_GLOBAL void pack_fp32_into_ue8m0(float* sf, uint32_t* out, uint32_t* gro } else { group_k = layout_value; } - const auto num_group_sf_rows = math::ceil_div(group_k, gran_k); + const auto aligned_group_k = kUsePsumLayout ? math::align(group_k, k_alignment) : group_k; + const auto num_group_sf_rows = aligned_group_k / gran_k; group_sf_row_start = num_prefix_sf_rows; group_sf_row_end = group_sf_row_start + num_group_sf_rows; num_prefix_sf_rows += num_group_sf_rows; @@ -225,14 +272,14 @@ CUTLASS_GLOBAL void pack_fp32_into_ue8m0(float* sf, uint32_t* out, uint32_t* gro #pragma unroll for (uint32_t j = 0; j < 4; ++ j) { values[j] = make_uint4(0, 0, 0, 0); - const uint32_t sf_row_idx = packed_sf_k_idx * 4 + j - num_padding_sf_rows; + const uint32_t sf_row_idx = static_cast(packed_sf_k_idx) * 4 + j - num_padding_sf_rows; if (sf_row_idx >= group_sf_row_start and sf_row_idx < group_sf_row_end) values[j] = reinterpret_cast(sf + sf_row_idx * mn)[mn_idx]; // FP32 SFs must have a zero sign and mantissa (only the exponent is packed) - DG_DEVICE_ASSERT((values[j].x & 0x807fffffu) == 0); - DG_DEVICE_ASSERT((values[j].y & 0x807fffffu) == 0); - DG_DEVICE_ASSERT((values[j].z & 0x807fffffu) == 0); - DG_DEVICE_ASSERT((values[j].w & 0x807fffffu) == 0); + DG_TRAP_ONLY_DEVICE_ASSERT((values[j].x & 0x807fffffu) == 0); + DG_TRAP_ONLY_DEVICE_ASSERT((values[j].y & 0x807fffffu) == 0); + DG_TRAP_ONLY_DEVICE_ASSERT((values[j].z & 0x807fffffu) == 0); + DG_TRAP_ONLY_DEVICE_ASSERT((values[j].w & 0x807fffffu) == 0); } // Pack and store diff --git a/deep_gemm/include/deep_gemm/layout/gemm.cuh b/deep_gemm/include/deep_gemm/layout/gemm.cuh new file mode 100644 index 0000000000..7cb36611ae --- /dev/null +++ b/deep_gemm/include/deep_gemm/layout/gemm.cuh @@ -0,0 +1,67 @@ +#pragma once + +#include +#include + +#include +#include + +namespace deep_gemm::layout { + +template +struct SM100BF16GemmSharedStorage { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using cd_dtype = cd_dtype_t; + using cd_stage_t = cd_dtype_t[STORE_BLOCK_M * STORE_BLOCK_N]; + using a_stage_t = cutlass::bfloat16_t[LOAD_BLOCK_M * BLOCK_K]; + using b_stage_t = cutlass::bfloat16_t[LOAD_BLOCK_N * BLOCK_K]; + DG_STATIC_ASSERT(kNumTMAStoreStages * sizeof(cd_stage_t) % 1024 == 0 and + sizeof(a_stage_t) % 1024 == 0 and sizeof(b_stage_t) % 1024 == 0, + "Shared memory of A/B must be aligned to 1024 bytes"); + alignas(1024) cd_stage_t cd[kNumTMAStoreStages]; + a_stage_t a[kNumStages]; + b_stage_t b[kNumStages]; + Barrier full_barriers[kNumStages]; + Barrier empty_barriers[kNumStages]; + Barrier tmem_full_barriers[kNumEpilogueStages]; + Barrier tmem_empty_barriers[kNumEpilogueStages]; + // Keep the BF16 barrier layout aligned with the block-scaled GEMM layout. + Barrier reserved_barriers[kNumStages]; + Barrier tensor_core_full_barrier; + uint32_t tmem_ptr; +}; + +template +struct SM100FP8FP4GemmSharedStorage { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using cd_dtype = cd_dtype_t; + using cd_stage_t = cd_dtype_t[STORE_BLOCK_M * STORE_BLOCK_N]; + using a_stage_t = a_dtype_t[LOAD_BLOCK_M * BLOCK_K / get_smem_pack_factor()]; + using b_stage_t = b_dtype_t[LOAD_BLOCK_N * BLOCK_K / get_smem_pack_factor()]; + using sfa_stage_t = uint32_t[SF_BLOCK_M * SF_BLOCK_K]; + using sfb_stage_t = uint32_t[SF_BLOCK_N * SF_BLOCK_K]; + DG_STATIC_ASSERT(kNumTMAStoreStages * sizeof(cd_stage_t) % 1024 == 0 and + sizeof(a_stage_t) % 1024 == 0 and sizeof(b_stage_t) % 1024 == 0, + "Shared memory of A/B must be aligned to 1024 bytes"); + alignas(1024) cd_stage_t cd[kNumTMAStoreStages]; + a_stage_t a[kNumStages]; + b_stage_t b[kNumStages]; + sfa_stage_t sfa[kNumStages]; + sfb_stage_t sfb[kNumStages]; + Barrier full_barriers[kNumStages]; + Barrier sf_full_barriers[kNumStages]; + Barrier empty_barriers[kNumStages]; + Barrier tmem_full_barriers[kNumEpilogueStages]; + Barrier tmem_empty_barriers[kNumEpilogueStages]; + Barrier tmem_overlap_barriers[kNumEpilogueStages]; + uint32_t tmem_ptr; +}; + +} // namespace deep_gemm::layout diff --git a/deep_gemm/include/deep_gemm/layout/mega_gate.cuh b/deep_gemm/include/deep_gemm/layout/mega_gate.cuh new file mode 100644 index 0000000000..05d1f05c4a --- /dev/null +++ b/deep_gemm/include/deep_gemm/layout/mega_gate.cuh @@ -0,0 +1,130 @@ +#pragma once + +#include +#include + +#include + +namespace deep_gemm::layout::mega_gate { + +// Kernel geometry and alignments +static constexpr uint32_t BLOCK_K = 64; +static constexpr uint32_t kNumNonEpilogueThreads = 128; +static constexpr uint32_t kNumEpilogueStages = 2; +static constexpr uint32_t kNumMaxBlockTokens = 256; +static constexpr uint32_t kExpertAlignment = 128; +static constexpr uint32_t kSharedMemoryAlignment = 1024; + +// Token bounds +static constexpr uint32_t kNumMinBlockTokens = 16; +static constexpr uint32_t kNumMaxTokens = 1u << 20; +static constexpr uint32_t kNumMaxTokenBlocks = math::constexpr_ceil_div(kNumMaxTokens, kNumMinBlockTokens); + +// Score-barrier workspace geometry +static constexpr uint32_t kScoreBarrierLineBytes = 128; +static constexpr uint32_t kNumMaxLogicalCtas = 64; +static constexpr uint32_t kNumMaxMetadataCacheBytes = 4 * 1024; + +enum class ScoringType : uint32_t { + Sigmoid = 0, + SqrtSoftplus = 1, + Identity = 3, +}; + +CUTLASS_HOST_DEVICE constexpr uint32_t get_num_bias_cache_bytes(const uint32_t num_routed_experts, + const bool has_bias, + const bool has_image_token_mask) { + return num_routed_experts * ((has_bias ? 1u : 0u) + (has_image_token_mask ? 1u : 0u)) * + static_cast(sizeof(float)); +} + +CUTLASS_HOST_DEVICE constexpr bool caches_logical_count(const uint32_t num_routed_experts, + const bool has_bias, const bool has_image_token_mask, + const bool to_physical_map_exists) { + return to_physical_map_exists and + get_num_bias_cache_bytes(num_routed_experts, has_bias, has_image_token_mask) + + num_routed_experts * static_cast(sizeof(int)) <= kNumMaxMetadataCacheBytes; +} + +CUTLASS_HOST_DEVICE constexpr uint32_t get_num_metadata_cache_bytes(const uint32_t num_routed_experts, + const bool has_bias, + const bool has_image_token_mask, + const bool to_physical_map_exists) { + return get_num_bias_cache_bytes(num_routed_experts, has_bias, has_image_token_mask) + + (caches_logical_count(num_routed_experts, has_bias, has_image_token_mask, to_physical_map_exists) + ? num_routed_experts * static_cast(sizeof(int)) : 0u); +} + +template +struct SharedStorage { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + using x_stage_t = cutlass::bfloat16_t[LOAD_BLOCK_M * BLOCK_K_]; + using weight_stage_t = cutlass::bfloat16_t[LOAD_BLOCK_N * BLOCK_K_]; + + static constexpr bool kCacheLogicalCount = caches_logical_count(kNumRoutedExperts, kHasBias, + kHasImageTokenMask, kToPhysicalMapExists); + static constexpr bool kHasMetadataCache = kHasBias or kHasImageTokenMask or kCacheLogicalCount; + static constexpr uint32_t kNumMetadataCacheElems = get_num_metadata_cache_bytes(kNumRoutedExperts, kHasBias, + kHasImageTokenMask, kToPhysicalMapExists) / sizeof(float); + + alignas(kSharedMemoryAlignment) x_stage_t x[kNumStages]; + alignas(kSharedMemoryAlignment) weight_stage_t weight[kNumStages]; + + alignas(Barrier) Barrier full_barriers[kNumStages]; + alignas(Barrier) Barrier empty_barriers[kNumStages]; + alignas(Barrier) Barrier tmem_full_barriers[kNumEpilogueStages]; + alignas(Barrier) Barrier tmem_empty_barriers[kNumEpilogueStages]; + alignas(Barrier) Barrier metadata_ready_barrier; + alignas(uint32_t) uint32_t tmem_ptr; + + alignas(4 * sizeof(float)) float metadata_cache[kNumMetadataCacheElems > 0 ? kNumMetadataCacheElems : 1]; + + CUTLASS_DEVICE float* get_bias_ptr() { + return metadata_cache; + } + + CUTLASS_DEVICE float* get_image_bias_ptr() { + return metadata_cache + (kHasBias ? kNumRoutedExperts : 0); + } + + CUTLASS_DEVICE int* get_logical_count_ptr() { + return reinterpret_cast(get_image_bias_ptr() + (kHasImageTokenMask ? kNumRoutedExperts : 0)); + } +}; + +// View over global score scratch: [num_token_blocks, kNumSplits, kBlockTokens, kNumExperts] +template +struct Workspace { + void* gmem_scratch; + void* gmem_score_barriers; + + CUTLASS_HOST_DEVICE + Workspace(void* gmem_scratch, void* gmem_score_barriers): + gmem_scratch(gmem_scratch), gmem_score_barriers(gmem_score_barriers) {} + + CUTLASS_HOST_DEVICE + static constexpr uint64_t get_num_scratch_bytes(const uint32_t num_token_blocks, + const uint32_t num_splits, + const uint32_t block_tokens, + const uint32_t num_experts) { + return static_cast(num_token_blocks) * num_splits * block_tokens * num_experts * sizeof(float); + } + + CUTLASS_HOST_DEVICE + float* get_score_ptr(const uint32_t token_block_idx, const uint32_t split_idx = 0) const { + return reinterpret_cast(gmem_scratch) + + (static_cast(token_block_idx) * kNumSplits + split_idx) * kBlockTokens * kNumExperts; + } + + CUTLASS_HOST_DEVICE uint64_t* get_score_barrier_ptr(const uint32_t token_block_idx) const { + return reinterpret_cast(reinterpret_cast(gmem_score_barriers) + + token_block_idx * kScoreBarrierLineBytes); + } +}; + +} // namespace deep_gemm::layout::mega_gate diff --git a/deep_gemm/include/deep_gemm/layout/mega_mhc.cuh b/deep_gemm/include/deep_gemm/layout/mega_mhc.cuh new file mode 100644 index 0000000000..de59a590e5 --- /dev/null +++ b/deep_gemm/include/deep_gemm/layout/mega_mhc.cuh @@ -0,0 +1,169 @@ +#pragma once + +#include + +#include + +namespace deep_gemm::layout::mega_mhc { + +// Model geometry +static constexpr uint32_t kNumRoutes = 4; +static constexpr uint32_t kNumHCOutputs = kNumRoutes * (kNumRoutes + 2); + +// FP8 scale-factor layout +static constexpr uint32_t kHiddenPerSF = 32; +static constexpr uint32_t kHiddenPerSFWord = 128; + +// Kernel tile and pipeline geometry +static constexpr uint32_t BLOCK_M = 64; +static constexpr uint32_t BLOCK_N = kNumHCOutputs; +static constexpr uint32_t BLOCK_K = 64; +static constexpr uint32_t kSwizzleMode = 128; +static constexpr uint32_t kFnAtomK = kSwizzleMode / sizeof(float); +static constexpr uint32_t kSwizzleAlignment = 8 * kSwizzleMode; +static constexpr uint32_t kNumIOStages = 4; +static constexpr uint32_t kNumFnStages = 2; +static constexpr uint32_t kNumAStages = 4; +static constexpr uint32_t kNumAccumStages = 2; +static constexpr uint32_t kNumCoeffStages = 2; +static constexpr uint32_t kNumWarpsPerWG = 4; +static constexpr uint32_t kNumThreads = 6 * kNumWarpsPerWG * 32; + +// Split-K workspace geometry +static constexpr uint32_t kNumGemmPartialElementsPerTask = BLOCK_M * kNumHCOutputs; +static constexpr uint32_t kNumSqrSumPartialElementsPerTask = BLOCK_M; +static constexpr uint32_t kDefaultNumSplits = 16; +static constexpr uint32_t kNumMaxSplits = 64; +static constexpr uint32_t kNumSplitBarriers = 2; +static constexpr uint64_t kSplitBarrierLineBytes = 128; +static constexpr uint32_t kNumMaxTokens = 1u << 20; +static constexpr uint32_t kNumMaxMBlocks = math::constexpr_ceil_div(kNumMaxTokens, BLOCK_M); + +// Host/device ABI payloads for the two consumer domains. Tensor maps remain +// standalone kernel parameters so TMA descriptors keep direct grid-constant access. +struct MixArgs { + const float* scales; + const float* bases; + float* new_prev_mix; + float* new_post_mix; + float* new_comb_res_mix; + float hc_norm_eps; + float hc_pre_eps; + float hc_post_scale; + float sinkhorn_eps; + uint32_t num_sinkhorn_iters; +}; + +struct NormArgs { + uint32_t num_tokens; + const nv_bfloat16* weight; + const nv_bfloat16* new_residual; + float eps; + float scale; + nv_bfloat16* y_bf16; + __nv_fp8_e4m3* y_fp8; + uint32_t* y_primary_sf; + int64_t y_primary_sf_stride_token; + int64_t y_primary_sf_stride_word; + uint32_t* y_shared_sf; + int64_t y_shared_sf_stride_word; +}; + +// View over global split partial scratch. +// +// Data layout: +// gemm_partials [num_m_blocks, kNumSplits, BLOCK_M, kNumHCOutputs] +// hc_norm_sqr_sum_partials [num_m_blocks, kNumSplits, BLOCK_M] +// x1_sqr_sum_partials [num_m_blocks, kNumSplits, BLOCK_M] +// where task_idx = m_block_idx * kNumSplits + k_split_idx. +template +struct Workspace { + void* gmem_scratch; + uint32_t num_m_blocks; + + DG_STATIC_ASSERT(kNumSplits > 0, "The number of splits must be positive"); + DG_STATIC_ASSERT(kNumSplits <= kNumMaxSplits, "The number of splits exceeds the supported maximum"); + + CUTLASS_HOST_DEVICE + Workspace(void* gmem_scratch, const uint32_t num_m_blocks): + gmem_scratch(gmem_scratch), num_m_blocks(num_m_blocks) {} + + CUTLASS_HOST_DEVICE + static constexpr uint32_t get_num_m_blocks(const uint32_t num_tokens) { + return math::constexpr_ceil_div(num_tokens, BLOCK_M); + } + + CUTLASS_HOST_DEVICE + static constexpr uint64_t get_num_scratch_bytes(const uint32_t num_m_blocks, const uint32_t num_splits = kNumSplits) { + const auto num_tasks = static_cast(num_m_blocks) * num_splits; + return num_tasks * (kNumGemmPartialElementsPerTask + 2 * kNumSqrSumPartialElementsPerTask) * sizeof(float); + } + + CUTLASS_HOST_DEVICE + float* get_gemm_partial_ptr(const uint32_t task_idx = 0) const { + return reinterpret_cast(gmem_scratch) + static_cast(task_idx) * kNumGemmPartialElementsPerTask; + } + + CUTLASS_HOST_DEVICE + float* get_hc_norm_sqr_sum_partial_ptr(const uint32_t task_idx = 0) const { + const auto base = get_gemm_partial_ptr(num_m_blocks * kNumSplits); + return base + static_cast(task_idx) * kNumSqrSumPartialElementsPerTask; + } + + CUTLASS_HOST_DEVICE + float* get_x1_sqr_sum_partial_ptr(const uint32_t task_idx = 0) const { + const auto base = get_hc_norm_sqr_sum_partial_ptr(num_m_blocks * kNumSplits); + return base + static_cast(task_idx) * kNumSqrSumPartialElementsPerTask; + } + +}; + +struct SharedStorage { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + // Route-major TMA layouts keep every 64 x 64 destination contiguous. + alignas(kSwizzleAlignment) nv_bfloat162 residual[kNumIOStages][kNumRoutes][BLOCK_M * BLOCK_K / 2]; + alignas(kSwizzleAlignment) nv_bfloat162 x[kNumIOStages][BLOCK_M * BLOCK_K / 2]; + + // Each Fn TMA fills all routes for one 32-element TF32 swizzle atom. + alignas(kSwizzleAlignment) float fn[kNumFnStages][BLOCK_K / kFnAtomK][kNumRoutes][kNumHCOutputs * kFnAtomK]; + + // Coeff tensors retain their TMA-native per-token layouts. + alignas(kSwizzleMode) float pre_coeff[kNumCoeffStages][BLOCK_M][kNumRoutes]; + alignas(kSwizzleMode) float post_coeff[kNumCoeffStages][BLOCK_M][kNumRoutes]; + alignas(kSwizzleMode) float comb_coeff[kNumCoeffStages][BLOCK_M][kNumRoutes * kNumRoutes]; + + // Post workers own independent partials reduced before the Mix/Norm barrier arrivals. + alignas(kSwizzleMode) float hc_norm_sqr_sums[2][BLOCK_M]; + alignas(kSwizzleMode) float x1_sqr_sums[2][BLOCK_M]; + // Fused Normal workers use dedicated scratch instead of aliasing Shifted statistics. + alignas(kSwizzleMode) float normal_scratch[BLOCK_M]; + + // Full/empty pairs connect static roles without CTA-wide block-boundary synchronization. + alignas(Barrier) Barrier full_io_barriers[kNumIOStages]; + alignas(Barrier) Barrier empty_io_barriers[kNumIOStages]; + alignas(Barrier) Barrier store_ready_io_barriers[kNumIOStages]; + + alignas(Barrier) Barrier full_fn_barriers[kNumFnStages]; + alignas(Barrier) Barrier empty_fn_barriers[kNumFnStages]; + + alignas(Barrier) Barrier full_a_barriers[kNumAStages]; + alignas(Barrier) Barrier empty_a_barriers[kNumAStages]; + + alignas(Barrier) Barrier full_accum_barriers[kNumAccumStages]; + alignas(Barrier) Barrier empty_accum_barriers[kNumAccumStages]; + + alignas(Barrier) Barrier full_coeff_barriers[kNumCoeffStages]; + alignas(Barrier) Barrier empty_coeff_barriers[kNumCoeffStages]; + + alignas(Barrier) Barrier full_stats_barrier; + alignas(Barrier) Barrier empty_stats_barrier; + + // Gate each split's Mix arrival on all WG3 warps and, in Normal, the output-store worker. + alignas(Barrier) Barrier mix_arrival_barrier; + alignas(uint32_t) uint32_t tmem_base; + // Shifted Norm and fused Normal workers share this mutually exclusive task queue. + alignas(uint32_t) uint32_t next_norm_task_ticket; +}; + +} // namespace deep_gemm::layout::mega_mhc diff --git a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh index 410c0b6a2a..8228a2b656 100644 --- a/deep_gemm/include/deep_gemm/layout/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/layout/mega_moe.cuh @@ -1,17 +1,19 @@ #pragma once +#include #include #include #include +#include namespace deep_gemm::layout { -static constexpr int kNumCandidateBlockMs = 7; -static constexpr int kCandidateBlockM[kNumCandidateBlockMs] = {8, 16, 32, 64, 96, 128, 192}; -static constexpr int kMaxCandidateBlockM = 192; +static constexpr int kNumCandidateBlockMs = 8; +static constexpr int kCandidateBlockM[kNumCandidateBlockMs] = {8, 16, 32, 64, 96, 128, 192, 240}; +static constexpr int kMaxCandidateBlockM = 240; static constexpr int kMinCandidateBlockM = 8; -static constexpr int kLCMCandidateBlockM = 384; +static constexpr int kLCMCandidateBlockM = 1920; // Pool capacity for shared expert token pool: worst-case total tokens + per-expert BLOCK_M alignment padding, among all possible BLOCK_M template @@ -41,27 +43,62 @@ struct TokenSrcMetadata { uint32_t rank_idx; uint32_t token_idx; uint32_t topk_idx; + + TokenSrcMetadata() = default; + + CUTLASS_DEVICE TokenSrcMetadata(const uint32_t& rank_idx, const uint32_t& token_idx, const uint32_t& topk_idx): + rank_idx(rank_idx), token_idx(token_idx), topk_idx(topk_idx) {} }; +struct alignas(128) MegaMoESignals { + static constexpr uint32_t kNumMaxGridSyncCounters = 4; + static constexpr uint32_t kNumMaxExperts = 2048; + static constexpr uint32_t kNumMaxExpertsPerRank = 512; + static constexpr uint32_t kNumMaxRingBlocks = 1u << 20; // 20 MiB for the ring signals + static constexpr uint32_t kNumMaxSharedL2Blocks = 1u << 15; // 128 KiB + + // Grid and NVLink synchronization + uint32_t grid_sync_count[kNumMaxGridSyncCounters]; + uint32_t nvl_barrier_counter; + int nvl_barrier_signals[2]; + + // Task scheduling + uint32_t l1_task_count; + uint32_t l2_task_count; + uint32_t shared_l1_task_count; + uint32_t shared_l2_task_count; + + // Combine readiness: `combine_ready_grid_idx[peer] == own grid index` means the peer's L2 writes into this + // rank are done; grid indices are unique per launch, so no reset is needed. Peers push theirs during dispatch. + alignas(128) uint64_t combine_ready_grid_idx[kNumMaxRanks]; + alignas(128) uint64_t peer_grid_idx[kNumMaxRanks]; + + // Expert token counts + alignas(128) uint64_t expert_send_count[kNumMaxExperts]; + uint64_t expert_recv_count[kNumMaxExperts]; + uint64_t expert_recv_count_sum[kNumMaxExpertsPerRank]; + + // Routed-expert ring signals, `l2_full_mask` has one bit per L1 N block + uint32_t l1_full_count[kNumMaxRingBlocks]; + uint32_t l1_empty_count[kNumMaxRingBlocks]; + uint64_t l2_full_mask[kNumMaxRingBlocks]; + uint32_t l2_empty_count[kNumMaxRingBlocks]; + + // Shared-expert signals + uint32_t shared_l2_full_count[kNumMaxSharedL2Blocks]; +}; +DG_STATIC_ASSERT(offsetof(MegaMoESignals, combine_ready_grid_idx) == 128, "Invalid control signal layout"); + struct Workspace { - void* base; + MegaMoESignals* signals; uint32_t num_ranks, num_experts; uint32_t num_experts_per_rank; uint32_t num_max_tokens_per_rank; - uint32_t num_max_recv_tokens_per_expert; - - // Ring-buffer capacity used by reusable token/data buffers - uint32_t num_ring_tokens; - uint32_t num_ring_blocks; uint32_t num_shared_l2_pool_blocks; // Full-pool span used by non-ring token metadata uint32_t num_max_pool_tokens; - // Keep grid/NVLink/schedule counters separated from expert counters. - // NVIDIA L2 cache lines are 128B, and these counters are hot atomics. - static constexpr uint64_t kNumBarrierSignalBytes = 128; - Workspace() = default; CUTLASS_HOST_DEVICE @@ -71,47 +108,31 @@ struct Workspace { const uint32_t& num_max_tokens_per_rank, const uint32_t& num_topk, const uint32_t& num_ring_tokens): - base(base), + signals(static_cast(base)), num_ranks(num_ranks), num_experts(num_experts), - num_max_tokens_per_rank(num_max_tokens_per_rank), - num_ring_tokens(num_ring_tokens) { + num_max_tokens_per_rank(num_max_tokens_per_rank) { num_experts_per_rank = num_experts / num_ranks; - num_max_recv_tokens_per_expert = num_ranks * num_max_tokens_per_rank; num_max_pool_tokens = get_num_max_pool_tokens(num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank); - num_ring_blocks = num_ring_tokens / kMinCandidateBlockM; num_shared_l2_pool_blocks = math::ceil_div(num_max_tokens_per_rank, kMinCandidateBlockM); + + DG_UNIFIED_ASSERT(num_ranks > 0); + DG_UNIFIED_ASSERT(num_ranks <= kNumMaxRanks); + DG_UNIFIED_ASSERT(num_experts % num_ranks == 0); + DG_UNIFIED_ASSERT(num_experts <= MegaMoESignals::kNumMaxExperts); + DG_UNIFIED_ASSERT(num_experts_per_rank <= MegaMoESignals::kNumMaxExpertsPerRank); + DG_UNIFIED_ASSERT(num_ring_tokens <= MegaMoESignals::kNumMaxRingBlocks * kMinCandidateBlockM); + DG_UNIFIED_ASSERT(num_shared_l2_pool_blocks <= MegaMoESignals::kNumMaxSharedL2Blocks); } CUTLASS_HOST_DEVICE uint64_t get_num_bytes() const { uint64_t num_bytes = 0; - // Barrier and in-kernel task scheduling counters - num_bytes += kNumBarrierSignalBytes; - - // Expert send/recv count - num_bytes += num_experts * sizeof(uint64_t) * 2; - - // Expert recv count sum - num_bytes += num_experts_per_rank * sizeof(uint64_t); - - // L1 full token count (ring) - num_bytes += num_ring_blocks * sizeof(uint32_t); - - // L1 empty block count (ring) - num_bytes += num_ring_blocks * sizeof(uint32_t); + // Fixed-position global and expert signals + num_bytes += sizeof(MegaMoESignals); - // L2 full block count (ring) - num_bytes += num_ring_blocks * sizeof(uint32_t); - - // L2 empty block count (ring) - num_bytes += num_ring_blocks * sizeof(uint32_t); - - // Shared L2 full block count - num_bytes += num_shared_l2_pool_blocks * sizeof(uint32_t); - - // Dispatch pulling source token-topk - num_bytes += num_experts_per_rank * num_ranks * num_max_recv_tokens_per_expert * sizeof(int); + // Source token-topk: [local expert][source rank][token], with distinct experts per token + num_bytes += num_experts * num_max_tokens_per_rank * sizeof(uint32_t); // Combine push source indices (full) num_bytes += num_max_pool_tokens * sizeof(TokenSrcMetadata); @@ -123,112 +144,104 @@ struct Workspace { CUTLASS_HOST_DEVICE void* get_end_ptr() const { - return math::advance_ptr(base, get_num_bytes()); + return math::advance_ptr(signals, get_num_bytes()); } - // Grid sync counters: `kNumBarrierSignalBytes` layout - // [ 0..15]: 4 x `uint32_t` grid sync counters - // [16..20]: `uint32_t` NVLink barrier counter - // [20..27]: 2 x `int` NVLink barrier signals (phase 0 and 1) - // [28..31]: `uint32_t` L1 schedule task counter - // [32..35]: `uint32_t` L2 schedule task counter - // [36..39]: `uint32_t` shared L1 schedule task counter - // [40..43]: `uint32_t` shared L2 schedule task counter - // [44..127]: padding to isolate hot expert counters from barrier/schedule counters - static constexpr uint32_t kNumMaxGridSyncCounters = 4; - template CUTLASS_DEVICE uint32_t* get_grid_sync_count_ptr() const { - DG_STATIC_ASSERT(kIndex < kNumMaxGridSyncCounters, "Grid sync index out of bounds"); - return static_cast(base) + kIndex; + DG_STATIC_ASSERT(kIndex < MegaMoESignals::kNumMaxGridSyncCounters, "Grid sync index out of bounds"); + return signals->grid_sync_count + kIndex; } CUTLASS_DEVICE uint32_t* get_nvl_barrier_counter_ptr() const { - return static_cast(base) + kNumMaxGridSyncCounters; + return &signals->nvl_barrier_counter; } CUTLASS_DEVICE int* get_nvl_barrier_signal_ptr(const uint32_t& phase) const { // NOTES: the signal is signed, as we may minus - return math::advance_ptr(base, (kNumMaxGridSyncCounters + 1) * sizeof(uint32_t) + phase * sizeof(int)); + return signals->nvl_barrier_signals + phase; + } + + CUTLASS_DEVICE + uint64_t* get_combine_ready_grid_idx_ptr(const uint32_t& peer_rank_idx = 0) const { + return signals->combine_ready_grid_idx + peer_rank_idx; + } + + CUTLASS_DEVICE + uint64_t* get_peer_grid_idx_ptr(const uint32_t& peer_rank_idx = 0) const { + return signals->peer_grid_idx + peer_rank_idx; } CUTLASS_DEVICE uint32_t* get_l1_task_count_ptr() const { - return math::advance_ptr(base, 28u); + return &signals->l1_task_count; } CUTLASS_DEVICE uint32_t* get_l2_task_count_ptr() const { - return math::advance_ptr(base, 32u); + return &signals->l2_task_count; } CUTLASS_DEVICE uint32_t* get_shared_l1_task_count_ptr() const { - return math::advance_ptr(base, 36u); + return &signals->shared_l1_task_count; } CUTLASS_DEVICE uint32_t* get_shared_l2_task_count_ptr() const { - return math::advance_ptr(base, 40u); + return &signals->shared_l2_task_count; } CUTLASS_DEVICE uint64_t* get_expert_send_count_ptr(const uint32_t& expert_idx = 0) const { - return math::advance_ptr(base, kNumBarrierSignalBytes) + expert_idx; + return signals->expert_send_count + expert_idx; } CUTLASS_DEVICE uint64_t* get_expert_recv_count_ptr( const uint32_t& rank_idx = 0, const uint32_t& expert_idx = 0) const { - return get_expert_send_count_ptr(num_experts) + rank_idx * num_experts_per_rank + expert_idx; + return signals->expert_recv_count + rank_idx * num_experts_per_rank + expert_idx; } CUTLASS_DEVICE uint64_t* get_expert_recv_count_sum_ptr(const uint32_t& expert_idx = 0) const { - return get_expert_send_count_ptr(num_experts * 2) + expert_idx; + return signals->expert_recv_count_sum + expert_idx; } CUTLASS_DEVICE uint32_t* get_l1_full_count_ptr(const uint32_t& ring_block_idx = 0) const { - const auto base = get_expert_recv_count_sum_ptr(num_experts_per_rank); - return reinterpret_cast(base) + ring_block_idx; + return signals->l1_full_count + ring_block_idx; } CUTLASS_DEVICE uint32_t* get_l1_empty_count_ptr(const uint32_t& ring_block_idx = 0) const { - const auto base = get_l1_full_count_ptr(num_ring_blocks); - return reinterpret_cast(base) + ring_block_idx; + return signals->l1_empty_count + ring_block_idx; } CUTLASS_DEVICE - uint32_t* get_l2_full_count_ptr(const uint32_t& ring_block_idx = 0) const { - const auto base = get_l1_empty_count_ptr(num_ring_blocks); - return reinterpret_cast(base) + ring_block_idx; + uint64_t* get_l2_full_mask_ptr(const uint32_t& ring_block_idx = 0) const { + return signals->l2_full_mask + ring_block_idx; } CUTLASS_DEVICE uint32_t* get_l2_empty_count_ptr(const uint32_t& ring_block_idx = 0) const { - const auto base = get_l2_full_count_ptr(num_ring_blocks); - return reinterpret_cast(base) + ring_block_idx; + return signals->l2_empty_count + ring_block_idx; } CUTLASS_DEVICE uint32_t* get_shared_l2_full_count_ptr(const uint32_t& block_idx = 0) const { - const auto base = get_l2_empty_count_ptr(num_ring_blocks); - return reinterpret_cast(base) + block_idx; + return signals->shared_l2_full_count + block_idx; } // For dispatch pulling CUTLASS_DEVICE uint32_t* get_src_token_topk_idx_ptr( const uint32_t& expert_idx = 0, const uint32_t& rank_idx = 0, const uint32_t& token_idx = 0) const { - const auto base = get_shared_l2_full_count_ptr(num_shared_l2_pool_blocks); - return reinterpret_cast(base) + - expert_idx * (num_ranks * num_max_recv_tokens_per_expert) + - rank_idx * num_max_recv_tokens_per_expert + token_idx; + const auto offset = (expert_idx * num_ranks + rank_idx) * num_max_tokens_per_rank + token_idx; + return math::advance_ptr(signals, sizeof(MegaMoESignals)) + offset; } // For combine usages (full) @@ -376,9 +389,9 @@ struct MegaMoEBuffer { const auto bf16_token_layout = layout::Data(hidden * 2); const auto intermediate_token_layout = layout::Data(intermediate_hidden * num_mma_elem_bytes); const auto shared_intermediate_token_layout = layout::Data(shared_intermediate_hidden * num_mma_elem_bytes); - const auto input_sf_layout = layout::Data(with_sf ? hidden / 32 : 0); - const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / 32 : 0); - const auto shared_intermediate_sf_layout = layout::Data(with_sf ? shared_intermediate_hidden / 32 : 0); + const auto input_sf_layout = layout::Data(with_sf ? hidden / 32 : 0, false); + const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / 32 : 0, false); + const auto shared_intermediate_sf_layout = layout::Data(with_sf ? shared_intermediate_hidden / 32 : 0, false); 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); @@ -438,7 +451,7 @@ struct MegaMoEBuffer { CUTLASS_HOST_DEVICE int64_t get_num_bytes() const { return static_cast(combine_token_buffer.get_end_ptr()) - - static_cast(workspace.base); + - reinterpret_cast(workspace.signals); } }; diff --git a/deep_gemm/include/deep_gemm/layout/mqa_logits.cuh b/deep_gemm/include/deep_gemm/layout/mqa_logits.cuh index 14485b54c1..cee77869aa 100644 --- a/deep_gemm/include/deep_gemm/layout/mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/layout/mqa_logits.cuh @@ -22,11 +22,13 @@ struct MQALogitsSharedStorage { static constexpr uint32_t kNumUTCCPAlignedElems = 128; static constexpr uint32_t kQKBytesPerElem = sizeof(qk_dtype_t); static constexpr uint32_t kNumQKBytesPerToken = kIsFP4 ? (kHeadDim / 2) : kHeadDim; + static constexpr uint32_t BLOCK_QH = BLOCK_Q * kNumHeads; + static constexpr uint32_t UMMA_N = math::constexpr_align(BLOCK_QH, 8u); // Align to one 8-row Q/K swizzle tile: FP4 uses head_dim / 2 bytes, FP8 use head_dim bytes per token static constexpr uint32_t kSwizzleAlignment = 8 * kNumQKBytesPerToken; - static constexpr uint32_t kNumSFQ = math::constexpr_align(BLOCK_Q * kNumHeads, kNumUTCCPAlignedElems); + static constexpr uint32_t kNumSFQ = math::constexpr_align(UMMA_N, kNumUTCCPAlignedElems); static constexpr uint32_t kNumSFKV = math::constexpr_align(SPLIT_KV, kNumUTCCPAlignedElems); - static constexpr uint32_t kNumQBytesPerStage = BLOCK_Q * kNumHeads * kNumQKBytesPerToken; + static constexpr uint32_t kNumQBytesPerStage = UMMA_N * kNumQKBytesPerToken; static constexpr uint32_t kNumKVBytesPerStage = SPLIT_KV * kNumQKBytesPerToken; static constexpr uint32_t kNumQElementsPerStage = kNumQBytesPerStage / kQKBytesPerElem; static constexpr uint32_t kNumKVElementsPerStage = kNumKVBytesPerStage / kQKBytesPerElem; @@ -35,17 +37,26 @@ struct MQALogitsSharedStorage { static constexpr uint32_t kNumScaleKV = kIsMXSF ? kNumSFKV : SPLIT_KV; // TMA destinations in shared memory must be 128-byte aligned. static constexpr uint32_t kTmaAlignment = 128; + static constexpr uint32_t kNumWeightBytesPerRow = math::constexpr_align( + kNumHeads * static_cast(sizeof(reduce_dtype_t)), 16u); + static constexpr uint32_t kNumWeightElementsPerRow = + kNumWeightBytesPerRow / static_cast(sizeof(reduce_dtype_t)); + static constexpr uint32_t kNumWeightBytesPerStage = math::constexpr_align( + BLOCK_Q * kNumWeightBytesPerRow, kTmaAlignment); + static constexpr uint32_t kNumWeightElementsPerStage = + kNumWeightBytesPerStage / static_cast(sizeof(reduce_dtype_t)); DG_STATIC_ASSERT(kNumQBytesPerStage % kSwizzleAlignment == 0, "Unaligned TMA swizzling"); DG_STATIC_ASSERT(kNumKVBytesPerStage % kSwizzleAlignment == 0, "Unaligned TMA swizzling"); DG_STATIC_ASSERT(kSwizzleAlignment % 128 == 0, "TMA destination must be 128-byte aligned"); DG_STATIC_ASSERT(kTmaAlignment % 128 == 0, "TMA destination must be 128-byte aligned"); + DG_STATIC_ASSERT(kNumWeightBytesPerStage % kTmaAlignment == 0, "Unaligned weight TMA stage"); alignas(kSwizzleAlignment) qk_dtype_t smem_q[kNumQStages][kNumQElementsPerStage]; alignas(kSwizzleAlignment) qk_dtype_t smem_kv[kNumKVStages][kNumKVElementsPerStage]; alignas(kTmaAlignment) sf_dtype_t smem_sf_q[kNumQStages][kNumScaleQ]; alignas(kTmaAlignment) sf_dtype_t smem_sf_kv[kNumKVStages][kNumScaleKV]; - alignas(kTmaAlignment) reduce_dtype_t smem_weights[kNumQStages][BLOCK_Q * kNumHeads]; + alignas(kTmaAlignment) reduce_dtype_t smem_weights[kNumQStages][kNumWeightElementsPerStage]; // Barriers require 8-byte alignment, already guaranteed by the preceding TMA-aligned arrays. Barrier full_q_barriers[kNumQStages]; Barrier empty_q_barriers[kNumQStages]; diff --git a/deep_gemm/include/deep_gemm/layout/sparse_mqa_logits.cuh b/deep_gemm/include/deep_gemm/layout/sparse_mqa_logits.cuh new file mode 100644 index 0000000000..7bb83f8884 --- /dev/null +++ b/deep_gemm/include/deep_gemm/layout/sparse_mqa_logits.cuh @@ -0,0 +1,146 @@ +#pragma once + +#include + +#include +#include +#include + +namespace deep_gemm::layout::sparse_mqa_logits { + +inline constexpr uint32_t kNumSparseSlotBits = 16; +inline constexpr uint32_t kInvalidSparseSlot = (1u << kNumSparseSlotBits) - 1; +inline constexpr uint32_t kNumHeads = 32; +inline constexpr uint32_t kHeadDim = 128; +inline constexpr uint32_t kNumUTCCPAlignedElems = 128; +inline constexpr uint32_t kNumKVTokensPerTMA = 128; +inline constexpr uint32_t kBlockQ = 2; + +CUTLASS_HOST_DEVICE constexpr uint32_t get_num_threads(const uint32_t num_math_warpgroups) { + return (num_math_warpgroups + 1 + num_math_warpgroups / 4) * 128; +} + +struct alignas(16) MetadataHeader { + uint32_t num_kv_splits; + uint32_t num_waves; + uint32_t use_unaligned_ks; +}; + +struct alignas(16) KVSplitHeader { + static constexpr uint32_t kContiguousFlag = 0x80000000u; + + uint32_t q_token_base; + uint32_t packed_num_kv_blocks; + uint32_t q0_slot_base; + uint32_t q1_slot_base; + + KVSplitHeader() = default; + CUTLASS_HOST_DEVICE KVSplitHeader(const uint32_t q_token_base, const uint32_t num_kv_blocks, + const bool is_contiguous, const uint32_t q0_slot_base, + const uint32_t q1_slot_base): + q_token_base(q_token_base), packed_num_kv_blocks(num_kv_blocks | (is_contiguous ? kContiguousFlag : 0)), + q0_slot_base(q0_slot_base), q1_slot_base(q1_slot_base) {} + + CUTLASS_HOST_DEVICE static constexpr uint32_t get_num_kv_blocks(const uint32_t packed_num_kv_blocks) { + return packed_num_kv_blocks & ~kContiguousFlag; + } + + CUTLASS_HOST_DEVICE static constexpr bool is_contiguous(const uint32_t packed_num_kv_blocks) { + return (packed_num_kv_blocks & kContiguousFlag) != 0; + } +}; + +struct KVBlockInfo { + uint32_t physical_kv_block_idx; + uint32_t packed_slot_offsets; + + KVBlockInfo() = default; + CUTLASS_HOST_DEVICE KVBlockInfo(const uint32_t physical_kv_block_idx, + const uint32_t q0_slot_offset, const uint32_t q1_slot_offset): + physical_kv_block_idx(physical_kv_block_idx), + packed_slot_offsets(q0_slot_offset | (q1_slot_offset << kNumSparseSlotBits)) {} +}; + +struct alignas(16) ScheduleEntry { + uint32_t kv_split_begin; + uint32_t kv_split_end; + uint32_t q_token_base; + uint32_t num_q_tokens; + + ScheduleEntry() = default; + CUTLASS_HOST_DEVICE ScheduleEntry(const uint32_t kv_split_begin, const uint32_t kv_split_end, + const uint32_t q_token_base, const uint32_t num_q_tokens): + kv_split_begin(kv_split_begin), kv_split_end(kv_split_end), q_token_base(q_token_base), num_q_tokens(num_q_tokens) {} +}; + +template +struct KVSplit { + KVSplitHeader header; + KVBlockInfo kv_block_infos[kNumKVBlocksPerSplit]; +}; + +struct QBlock { + uint32_t q_token_base; + uint32_t num_q_tokens; + uint32_t num_kv_splits; + + QBlock() = default; + CUTLASS_HOST_DEVICE QBlock(const uint32_t q_token_base, const uint32_t num_q_tokens, + const uint32_t num_kv_splits) + : q_token_base(q_token_base), num_q_tokens(num_q_tokens), num_kv_splits(num_kv_splits) {} +}; + +template +struct SharedStorage { + using Barrier = cutlass::arch::ClusterTransactionBarrier; + + static constexpr uint32_t kPackFactor = get_smem_pack_factor(); + static constexpr uint32_t kNumKVBlocksPerSplit = SPLIT_KV / SPARSE_BLOCK_KV; + static constexpr uint32_t kNumSFQ = math::constexpr_align(BLOCK_Q * kNumHeads, kNumUTCCPAlignedElems); + static constexpr uint32_t kSwizzleAlignment = 8 * kHeadDim / kPackFactor; + + DG_STATIC_ASSERT(BLOCK_Q == 2, "Sparse metadata packs exactly two Q slots"); + DG_STATIC_ASSERT(SPARSE_BLOCK_KV == 8 or SPARSE_BLOCK_KV == 16, "Invalid sparse KV block size"); + DG_STATIC_ASSERT(SPLIT_KV % SPARSE_BLOCK_KV == 0 and SPLIT_KV % 128 == 0, "Invalid sparse KV split size"); + DG_STATIC_ASSERT(kNumKVBlocksPerSplit <= kInvalidSparseSlot, + "Sparse split-local slots must not use the invalid-slot value"); + + alignas(kSwizzleAlignment) qk_dtype_t q[kNumQStages][BLOCK_Q * kNumHeads][kHeadDim / kPackFactor]; + alignas(kSwizzleAlignment) qk_dtype_t kv[kNumKVStages][SPLIT_KV][kHeadDim / kPackFactor]; + alignas(128) uint32_t sf_q[kNumQStages][kNumSFQ]; + alignas(128) uint32_t sf_kv[kNumKVStages][SPLIT_KV]; + alignas(128) nv_bfloat16 weights[kNumQStages][BLOCK_Q * kNumHeads]; + alignas(16) KVBlockInfo kv_block_infos[kNumKVStages][kNumKVBlocksPerSplit]; + alignas(16) QBlock q_blocks[kNumQStages]; + alignas(16) KVSplitHeader kv_split_headers[kNumKVStages]; + + Barrier full_q_barriers[kNumQStages]; + Barrier full_sf_q_barriers[kNumQStages]; + Barrier empty_q_barriers[kNumQStages]; + Barrier full_metadata_barriers[kNumKVStages]; + Barrier full_sf_copy_barriers[kNumKVStages]; + Barrier full_kv_barriers[kNumKVStages]; + Barrier empty_kv_barriers[kNumKVStages]; + Barrier full_tmem_barriers[kNumTmemStages]; + Barrier empty_tmem_barriers[kNumTmemStages]; + uint32_t tmem_ptr_in_smem; +}; + +// Metadata kernel workspace +struct alignas(128) WorkspaceState { + // Keep concurrently updated counters on separate L2 lines + uint32_t num_kv_splits; + alignas(128) uint32_t next_q_offset; + alignas(128) uint32_t num_finished_ctas; +}; + +struct alignas(8) QBlockInfo { + uint32_t kv_split_base; + uint32_t num_kv_splits; + + CUTLASS_HOST_DEVICE QBlockInfo(const uint32_t kv_split_base, const uint32_t num_kv_splits): + kv_split_base(kv_split_base), num_kv_splits(num_kv_splits) {} +}; + +} // namespace deep_gemm::layout::sparse_mqa_logits diff --git a/deep_gemm/include/deep_gemm/mma/sm100.cuh b/deep_gemm/include/deep_gemm/mma/sm100.cuh index 86b24fe418..f21a09eb7c 100644 --- a/deep_gemm/include/deep_gemm/mma/sm100.cuh +++ b/deep_gemm/include/deep_gemm/mma/sm100.cuh @@ -5,6 +5,7 @@ #include #include +#include #include namespace deep_gemm::mma::sm100 { @@ -82,24 +83,16 @@ constexpr static cute::UMMA::LayoutType to_umma_layout_type() { template CUTLASS_DEVICE constexpr uint32_t get_umma_desc_stride_k() { - return kMajorMode == cute::UMMA::Major::K ? 1 : tma::get_inner_block_atom_size(); + return kMajorMode == cute::UMMA::Major::K ? 1 : tma::get_inner_block_atom_size(); } template CUTLASS_DEVICE uint32_t advance_umma_desc_lo(const uint32_t& base, const uint32_t& offset, const uint32_t& k_idx) { - return base + (((offset + k_idx * get_umma_desc_stride_k()) * static_cast(sizeof(dtype_t))) >> 4u); -} - -template -CUTLASS_DEVICE -constexpr uint32_t get_umma_desc_pack_factor() { - // Packed FP4 stores two logical elements per byte in SMEM. - if constexpr (cute::is_same_v) { - return 2; - } else { - return 1; - } + constexpr uint32_t kPackFactor = get_smem_pack_factor(); + const uint32_t byte_offset = (offset + k_idx * get_umma_desc_stride_k()) + * static_cast(sizeof(dtype_t)) / kPackFactor; + return base + (byte_offset >> 4u); } template (); + constexpr uint32_t kPackFactor = get_smem_pack_factor(); DG_STATIC_ASSERT(kPackFactor == 1 or sizeof(dtype_t) == 1, "Packing expects a 1-byte storage type"); const uint32_t stride_k = get_umma_desc_stride_k(); const auto layout_type = to_umma_layout_type(); @@ -127,7 +120,7 @@ cute::UMMA::SmemDescriptor make_umma_desc(dtype_t* base_smem_ptr, uint32_t mn_id return make_smem_desc(layout_type, byte_ptr, stride_byte_offset, leading_byte_offset); } else { DG_STATIC_ASSERT(kPackFactor <= 1, "Packing only supports K-major"); - constexpr uint32_t BLOCK_MN_ATOM = tma::get_inner_block_atom_size(); + constexpr uint32_t BLOCK_MN_ATOM = tma::get_inner_block_atom_size(); // Must have no in-atom MN-idx // NOTES: no worries for the runtime assert, the `mn_idx` are constants at compilation time diff --git a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh index 806a4c2e20..5cfec6aaab 100644 --- a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh +++ b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh @@ -78,7 +78,7 @@ struct SM90_U32x4_STSM_T { template struct SM100_U8x4_STSM_T { - __device__ __forceinline__ static void + CUTLASS_DEVICE static void copy(dtype_t src_0, void* smem_dst) { DG_STATIC_ASSERT(sizeof(dtype_t) == sizeof(uint32_t), "Invalid dtype"); const uint32_t src = *reinterpret_cast(&src_0); @@ -89,7 +89,7 @@ struct SM100_U8x4_STSM_T { template struct SM100_U8x8_STSM_T { - __device__ __forceinline__ static void + CUTLASS_DEVICE static void copy(dtype_t src_0, dtype_t src_1, void* smem_dst) { DG_STATIC_ASSERT(sizeof(dtype_t) == sizeof(uint32_t), "Invalid dtype"); const uint32_t src[2] = {*reinterpret_cast(&src_0), *reinterpret_cast(&src_1)}; @@ -153,6 +153,18 @@ CUTLASS_DEVICE void st_shared(const __int128_t* ptr, __int128_t val) { asm volatile("st.shared.b128 [%0], %1;" :: "l"(__cvta_generic_to_shared(ptr)), "q"(val)); } +template +CUTLASS_DEVICE void cp_async_cg(const uint4* global_ptr, uint4* shared_ptr) { + DG_STATIC_ASSERT(kNumL2PrefetchBytes == 64 or kNumL2PrefetchBytes == 256, "Invalid L2 prefetch size"); + if constexpr (kNumL2PrefetchBytes == 64) { + asm volatile("cp.async.cg.shared::cta.global.L2::64B [%0], [%1], 16;" :: + "r"(static_cast(__cvta_generic_to_shared(shared_ptr))), "l"(global_ptr)); + } else { + asm volatile("cp.async.cg.shared::cta.global.L2::256B [%0], [%1], 16;" :: + "r"(static_cast(__cvta_generic_to_shared(shared_ptr))), "l"(global_ptr)); + } +} + CUTLASS_DEVICE uint32_t mapa_shared(const uint32_t& ptr, const uint32_t& dst_cta_idx) { uint32_t mapped; asm volatile("mapa.shared::cluster.u32 %0, %1, %2;" : "=r"(mapped) : "r"(ptr), "r"(dst_cta_idx)); @@ -187,6 +199,34 @@ CUTLASS_DEVICE void st_shared_bulk(void* smem_ptr, const uint32_t& num_bytes) { } /// Global memory +constexpr uint64_t kEvictFirstCacheHint = 0x12f0000000000000; + +CUTLASS_DEVICE uint4 ld_evict_first(const uint4* ptr) { + uint4 value; + asm volatile("ld.weak.global.L1::no_allocate.L2::cache_hint.v4.b32 {%0, %1, %2, %3}, [%4], %5;" + : "=r"(value.x), "=r"(value.y), "=r"(value.z), "=r"(value.w) + : "l"(ptr), "l"(kEvictFirstCacheHint)); + return value; +} + +CUTLASS_DEVICE float ld_global(const float* ptr) { + float ret; + asm volatile("ld.weak.global.f32 %0, [%1];" : "=f"(ret) : "l"(ptr) : "memory"); + return ret; +} + +CUTLASS_DEVICE float4 ld_global(const float4* ptr) { + float4 ret; + asm volatile("ld.weak.global.v4.f32 {%0, %1, %2, %3}, [%4];" + : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w) + : "l"(ptr) : "memory"); + return ret; +} + +CUTLASS_DEVICE void st_global(float* ptr, const float& value) { + asm volatile("st.weak.global.f32 [%0], %1;" :: "l"(ptr), "f"(value) : "memory"); +} + CUTLASS_DEVICE uint64_t ld_volatile(const uint64_t* ptr) { uint64_t ret; asm volatile("ld.volatile.global.b64 %0, [%1];" : "=l"(ret) : "l"(ptr)); @@ -201,13 +241,13 @@ CUTLASS_DEVICE uint32_t ld_volatile(const uint32_t* ptr) { CUTLASS_DEVICE uint32_t ld_acq(const uint32_t* ptr) { uint32_t ret; - asm volatile("ld.acquire.gpu.global.b32 %0, [%1];" : "=r"(ret) : "l"(ptr)); + asm volatile("ld.acquire.gpu.global.b32 %0, [%1];" : "=r"(ret) : "l"(ptr) : "memory"); return ret; } CUTLASS_DEVICE uint64_t ld_acq_sys(const uint64_t* ptr) { uint64_t ret; - asm volatile("ld.acquire.sys.global.b64 %0, [%1];" : "=l"(ret) : "l"(ptr)); + asm volatile("ld.acquire.sys.global.b64 %0, [%1];" : "=l"(ret) : "l"(ptr) : "memory"); return ret; } @@ -236,10 +276,14 @@ CUTLASS_DEVICE uint64_t atomic_add_sys(const uint64_t* ptr, const uint64_t& valu CUTLASS_DEVICE uint32_t atomic_add_rel(const uint32_t* ptr, const uint32_t& value) { uint32_t ret; - asm volatile("atom.release.gpu.global.add.u32 %0, [%1], %2;" : "=r"(ret) : "l"(ptr), "r"(value)); + asm volatile("atom.release.gpu.global.add.u32 %0, [%1], %2;" : "=r"(ret) : "l"(ptr), "r"(value) : "memory"); return ret; } +CUTLASS_DEVICE void st_rel(const uint64_t* ptr, const uint64_t value) { + asm volatile("st.release.gpu.global.u64 [%0], %1;" :: "l"(ptr), "l"(value) : "memory"); +} + CUTLASS_DEVICE void red_add(const int* ptr, const int& value) { asm volatile("red.gpu.global.add.s32 [%0], %1;" :: "l"(ptr), "r"(value)); } @@ -256,10 +300,22 @@ CUTLASS_DEVICE void red_or_rel_gpu(uint64_t* ptr, const uint64_t& value) { asm volatile("red.release.gpu.global.or.b64 [%0], %1;" :: "l"(ptr), "l"(value)); } +CUTLASS_DEVICE void red_xor_rel(const uint64_t* ptr, const uint64_t& value) { + asm volatile("red.release.gpu.global.xor.b64 [%0], %1;" :: "l"(ptr), "l"(value) : "memory"); +} + +CUTLASS_DEVICE void st_rel_sys(const uint64_t* ptr, const uint64_t& value) { + asm volatile("st.release.sys.global.u64 [%0], %1;" :: "l"(ptr), "l"(value) : "memory"); +} + CUTLASS_DEVICE void red_add_rel(const uint32_t* ptr, const uint32_t& value) { asm volatile("red.release.gpu.global.add.u32 [%0], %1;" :: "l"(ptr), "r"(value)); } +CUTLASS_DEVICE void red_async_inc_rel(const uint64_t* ptr) { + asm volatile("red.async.release.gpu.global.add.u64 [%0], 1;" :: "l"(ptr) : "memory"); +} + CUTLASS_DEVICE void red_add_rel_sys(const int* ptr, const int& value) { asm volatile("red.release.sys.global.add.s32 [%0], %1;" :: "l"(ptr), "r"(value)); } @@ -278,7 +334,7 @@ CUTLASS_DEVICE uint32_t ld_acq_sys(const uint32_t* ptr) { CUTLASS_DEVICE uint64_t ld_acq_gpu(const uint64_t* ptr) { uint64_t ret; - asm volatile("ld.acquire.gpu.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr)); + asm volatile("ld.acquire.gpu.global.u64 %0, [%1];" : "=l"(ret) : "l"(ptr) : "memory"); return ret; } diff --git a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh index 528b3dd103..48426dd9da 100644 --- a/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh +++ b/deep_gemm/include/deep_gemm/ptx/tcgen05.cuh @@ -1,5 +1,11 @@ #pragma once +#include + +#include + +#include + namespace deep_gemm::ptx { /// UMMA versions with relaxed assertions @@ -139,6 +145,30 @@ struct SM100_MMA_MXF4_SS { } }; +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, @@ -157,6 +187,27 @@ struct SM100_MMA_F16BF16_WS_SS { }; /// Tensor memory operations +template +CUTLASS_DEVICE void tmem_load_32dp32b(const uint32_t tmem_addr, uint32_t* values) { + DG_STATIC_ASSERT(kNumValues == 4 or kNumValues == 8 or kNumValues == 16 or + kNumValues == 32 or kNumValues == 64 or kNumValues == 128, + "Invalid TMEM load width"); + using Loader = cute::conditional_t>>>>; + [&](cute::index_sequence) { + Loader::copy(tmem_addr, values[Is]...); + }(cute::make_index_sequence{}); +} + +CUTLASS_DEVICE void umma_arrive_no_elect(cutlass::arch::ClusterTransactionBarrier& barrier) { + asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];" :: + "r"(static_cast(__cvta_generic_to_shared(&barrier)))); +} + CUTLASS_DEVICE void tcgen05_before_thread_sync() { asm volatile("tcgen05.fence::before_thread_sync;"); } diff --git a/deep_gemm/include/deep_gemm/ptx/tma.cuh b/deep_gemm/include/deep_gemm/ptx/tma.cuh index 1530a3edc5..af04b66bec 100644 --- a/deep_gemm/include/deep_gemm/ptx/tma.cuh +++ b/deep_gemm/include/deep_gemm/ptx/tma.cuh @@ -38,6 +38,23 @@ CUTLASS_DEVICE void mbarrier_arrive( "r"(static_cast(__cvta_generic_to_shared(ptr)))); } +CUTLASS_DEVICE void mbarrier_arrive_count( + cutlass::arch::ClusterTransactionBarrier& barrier, const uint32_t count) { + asm volatile("mbarrier.arrive.shared::cta.b64 _, [%0], %1;" :: + "r"(static_cast(__cvta_generic_to_shared(&barrier))), "r"(count)); +} + +CUTLASS_DEVICE void mbarrier_arrive_count_pred( + cutlass::arch::ClusterTransactionBarrier& barrier, const uint32_t count, const bool pred) { + asm volatile( + "{\n\t" + ".reg .pred p;\n\t" + "setp.ne.b32 p, %2, 0;\n\t" + "@p mbarrier.arrive.shared::cta.b64 _, [%0], %1;\n\t" + "}" :: + "r"(static_cast(__cvta_generic_to_shared(&barrier))), "r"(count), "r"(static_cast(pred))); +} + CUTLASS_DEVICE void mbarrier_arrive_and_set_tx( cutlass::arch::ClusterTransactionBarrier* ptr, const uint32_t& num_bytes) { asm volatile("mbarrier.arrive.expect_tx.shared::cta.b64 _, [%1], %0; \n\t" :: @@ -64,7 +81,7 @@ CUTLASS_DEVICE void tma_load_1d( const void* dst_ptr, const void* src_ptr, cutlass::arch::ClusterTransactionBarrier* mbarrier_ptr, const uint32_t& num_bytes, - const cute::TMA::CacheHintSm90& hint = cute::TMA::CacheHintSm90::EVICT_FIRST) { + const cute::TMA::CacheHintSm100& hint = cute::TMA::CacheHintSm100::EVICT_FIRST) { // NOTES: normally, the loaded part will be evicted soon asm volatile( "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes.L2::cache_hint [%0], [%1], %2, [%3], %4;\n" :: @@ -78,7 +95,7 @@ CUTLASS_DEVICE void tma_load_1d( CUTLASS_DEVICE void tma_store_1d( const void* dst_ptr, const void* src_ptr, const uint32_t& num_bytes, - const cute::TMA::CacheHintSm90& hint = cute::TMA::CacheHintSm90::EVICT_NORMAL) { + const cute::TMA::CacheHintSm100& hint = cute::TMA::CacheHintSm100::EVICT_NORMAL) { // NOTES: normally, the stored part will be used soon asm volatile("cp.async.bulk.global.shared::cta.bulk_group.L2::cache_hint [%0], [%1], %2, %3;\n" :: "l"(dst_ptr), @@ -88,6 +105,32 @@ CUTLASS_DEVICE void tma_store_1d( : "memory"); } +CUTLASS_DEVICE void tma_store_2d( + const void* desc_ptr, const void* smem_ptr, + const uint32_t& crd_0, const uint32_t& crd_1, + const cute::TMA::CacheHintSm100& hint = cute::TMA::CacheHintSm100::EVICT_NORMAL) { + const auto gmem_int_desc = reinterpret_cast(desc_ptr); + const auto smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "cp.async.bulk.tensor.2d.global.shared::cta.bulk_group.L2::cache_hint " + "[%0, {%2, %3}], [%1], %4;\n" :: + "l"(gmem_int_desc), "r"(smem_int_ptr), + "r"(crd_0), "r"(crd_1), "l"(hint) + : "memory"); +} + +CUTLASS_DEVICE void tma_store_3d( + const void* desc_ptr, const void* smem_ptr, + const uint32_t& crd_0, const uint32_t& crd_1, const uint32_t& crd_2, + const cute::TMA::CacheHintSm100& hint) { + const auto gmem_int_desc = reinterpret_cast(desc_ptr); + const auto smem_int_ptr = static_cast(__cvta_generic_to_shared(smem_ptr)); + asm volatile( + "cp.async.bulk.tensor.3d.global.shared::cta.bulk_group.L2::cache_hint " + "[%0, {%2, %3, %4}], [%1], %5;\n" :: + "l"(gmem_int_desc), "r"(smem_int_ptr), "r"(crd_0), "r"(crd_1), "r"(crd_2), "l"(hint) : "memory"); +} + template __forceinline__ __device__ void tma_store_wait() { // NOTES: this function does not have `.read` diff --git a/deep_gemm/include/deep_gemm/ptx/utils.cuh b/deep_gemm/include/deep_gemm/ptx/utils.cuh index 16607dca6c..5ef3f46e22 100644 --- a/deep_gemm/include/deep_gemm/ptx/utils.cuh +++ b/deep_gemm/include/deep_gemm/ptx/utils.cuh @@ -20,6 +20,12 @@ CUTLASS_DEVICE uint32_t get_lane_idx() { return lane_id; } +CUTLASS_DEVICE uint64_t get_grid_idx() { + uint64_t grid_idx; + asm volatile("mov.u64 %0, %%gridid;" : "=l"(grid_idx)); + return grid_idx; +} + CUTLASS_DEVICE void sync_aligned(const uint32_t& num_threads, const uint32_t& barrier_idx) { asm volatile("bar.sync %0, %1;" : : "r"(barrier_idx), "r"(num_threads)); } @@ -40,6 +46,12 @@ CUTLASS_DEVICE dtype_t exchange(dtype_t ptr, const uint32_t& src_lane_idx) { return recv_dtype; } +CUTLASS_DEVICE float reduce_max_sync(const float& value) { + float ret; + asm volatile("redux.sync.max.f32 %0, %1, 0xffffffff;" : "=f"(ret) : "f"(value)); + return ret; +} + CUTLASS_DEVICE nv_bfloat162 cvt_relu_bf16x2_f32(const float2& v) { #if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) uint32_t packed; @@ -61,5 +73,17 @@ CUTLASS_DEVICE void accumulate(float2& a, nv_bfloat162 b) { #endif } +CUTLASS_DEVICE void accumulate_square(float2& a, nv_bfloat162 b) { +#if defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 1000) + // Accumulate BF16 squares directly into FP32 without separate conversion instructions. + asm("fma.rn.f32.bf16 %0, %1, %1, %0;\n" : "+f"(a.x) : "h"(*reinterpret_cast(&b.x))); + asm("fma.rn.f32.bf16 %0, %1, %1, %0;\n" : "+f"(a.y) : "h"(*reinterpret_cast(&b.y))); +#else + const float2 b_f32 = __bfloat1622float2(b); + a.x += b_f32.x * b_f32.x; + a.y += b_f32.y * b_f32.y; +#endif +} + } // namespace deep_gemm::ptx #endif diff --git a/deep_gemm/include/deep_gemm/scheduler/gemm.cuh b/deep_gemm/include/deep_gemm/scheduler/gemm.cuh index dbfb28bf20..23af6d39b1 100644 --- a/deep_gemm/include/deep_gemm/scheduler/gemm.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/gemm.cuh @@ -34,9 +34,14 @@ template ()> struct Scheduler { + // A/B group starts must be aligned to whole K blocks. SF rows are packed + // independently per group and tracked by `current_sf_k_cumsum`. + DG_STATIC_ASSERT(not is_k_grouped_contiguous(kGemmType) or kKAlignment % 128 == 0, + "K alignment must be a multiple of BLOCK_K (128)"); + int current_iter = -1; // Block configs @@ -56,31 +61,23 @@ struct Scheduler { // Only used for contiguous psum layout uint32_t last_psum_m = 0, current_psum_m, current_m_block_cumsum = 0; // Only used for k-grouped layout - uint32_t current_shape_k, current_num_valid_groups = 0, current_k_cumsum = 0, current_sf_k_cumsum = 0; - // NOTES: only used by the non-psum path; the psum path never reads them. - uint32_t next_group_idx, next_shape_k; + // NOTES: `current_k_start` is the current group's physical K start offset + // (always a multiple of `kKAlignment`), maintained by both psum and non-psum paths + uint32_t current_shape_k, current_k_start = 0, current_sf_k_cumsum = 0; // Only used for `KGroupedContiguousWithPsumLayout` - uint32_t current_k_start = 0, current_k_end = 0; - - // Only used for k-grouped gemm - CUTLASS_DEVICE void get_next_k_group(uint32_t &group_idx, uint32_t &shape_k) const { - for (; group_idx < kNumGroups; ++ group_idx) { - shape_k = grouped_layout[group_idx]; - if (shape_k > 0) - break; - } - } - - CUTLASS_DEVICE void get_next_psum_k_group(uint32_t &group_idx, uint32_t &shape_k, - uint32_t &k_start, uint32_t &k_end) const { - // NOTES: `grouped_layout[i]` is the psum end offset (K elements); each group starts at `align(prev_end, kKAlignment)`. Skip empty groups. - for (; group_idx < kNumGroups; ++ group_idx) { - const auto next_k_end = static_cast(grouped_layout[group_idx]); - k_start = math::align(k_end, kKAlignment); - shape_k = next_k_end - k_start; - k_end = next_k_end; - if (shape_k > 0) - break; + uint32_t current_k_end = 0; + + // Load the K-group selected by `current_group_idx`. + CUTLASS_DEVICE void get_next_k_group() { + if constexpr (kGemmType == GemmType::KGroupedContiguousWithPsumLayout) { + // `grouped_layout[i]` is the psum end offset in K elements. + const auto next_k_end = static_cast(grouped_layout[current_group_idx]); + current_k_start = math::align(current_k_end, kKAlignment); + current_shape_k = next_k_end - current_k_start; + current_k_end = next_k_end; + } else { + current_k_start += current_shape_k; + current_shape_k = grouped_layout[current_group_idx]; } } @@ -89,7 +86,7 @@ struct Scheduler { const uint32_t& shape_k, int* grouped_layout = nullptr) { num_m_blocks = math::ceil_div(shape_m, BLOCK_M); num_n_blocks = math::ceil_div(shape_n, BLOCK_N); - current_shape_k = shape_k; + current_shape_k = is_k_grouped_contiguous(kGemmType) ? 0 : shape_k; if constexpr (kGemmType == GemmType::Normal or kGemmType == GemmType::Batched) { num_blocks = num_m_blocks * num_n_blocks; } else if constexpr (kGemmType == GemmType::MGroupedContiguous) { @@ -104,13 +101,7 @@ struct Scheduler { } else if constexpr (is_k_grouped_contiguous(kGemmType)) { num_blocks = num_m_blocks * num_n_blocks; this->grouped_layout = grouped_layout; - if constexpr (kGemmType == GemmType::KGroupedContiguousWithPsumLayout) { - get_next_psum_k_group(current_group_idx, current_shape_k, current_k_start, current_k_end); - } else { - get_next_k_group(current_group_idx, current_shape_k); - next_group_idx = current_group_idx + 1; - get_next_k_group(next_group_idx, next_shape_k); - } + get_next_k_group(); } } @@ -169,10 +160,7 @@ struct Scheduler { if constexpr (kIndexType == IndexType::MN) { offset = current_group_idx * shape_dim; } else if constexpr (kIndexType == IndexType::K) { - if constexpr (kGemmType == GemmType::KGroupedContiguousWithPsumLayout) - offset = current_k_start; - else - offset = current_k_cumsum; + offset = current_k_start; } else if constexpr (kIndexType == IndexType::SF_K) { offset = current_sf_k_cumsum; } @@ -242,23 +230,20 @@ struct Scheduler { return false; // Within current group - if (next_block_idx < (current_num_valid_groups + 1) * num_blocks) + if (next_block_idx < (current_group_idx + 1) * num_blocks) break; // Move to check the next group - current_sf_k_cumsum += math::ceil_div(current_shape_k, kSFKSpan); - current_num_valid_groups ++; - if constexpr (kGemmType == GemmType::KGroupedContiguousWithPsumLayout) { - get_next_psum_k_group(++ current_group_idx, current_shape_k, current_k_start, current_k_end); - } else { - current_k_cumsum += current_shape_k; - current_group_idx = next_group_idx ++; - current_shape_k = next_shape_k; - get_next_k_group(next_group_idx, next_shape_k); - } + current_group_idx ++; + if (current_group_idx >= kNumGroups) + return false; + + const auto aligned_shape_k = math::align(current_shape_k, kKAlignment); + current_sf_k_cumsum += math::ceil_div(aligned_shape_k, kSFKSpan); + get_next_k_group(); } - get_swizzled_block_idx(next_block_idx - current_num_valid_groups * num_blocks, m_block_idx, n_block_idx); + get_swizzled_block_idx(next_block_idx - current_group_idx * num_blocks, m_block_idx, n_block_idx); } else if constexpr (kGemmType == GemmType::Batched) { if (next_block_idx >= num_blocks * kNumGroups) return false; diff --git a/deep_gemm/include/deep_gemm/scheduler/mega_gate.cuh b/deep_gemm/include/deep_gemm/scheduler/mega_gate.cuh new file mode 100644 index 0000000000..d400bdc359 --- /dev/null +++ b/deep_gemm/include/deep_gemm/scheduler/mega_gate.cuh @@ -0,0 +1,55 @@ +#pragma once + +#include +#include +#include + +namespace deep_gemm::sched::mega_gate { + +template +struct Scheduler { + static constexpr uint32_t kNumWorkerGroups = kNumSMs / kNumLogicalCtas; + uint32_t worker_group_idx; + uint32_t num_token_blocks; + + CUTLASS_DEVICE Scheduler(const uint32_t& num_tokens): + worker_group_idx(blockIdx.x / kNumLogicalCtas), + num_token_blocks(math::ceil_div(num_tokens, kBlockTokens)) {} + + template + CUTLASS_DEVICE void run(const Task& run_task) const { + uint32_t iter_idx = 0; + for (auto token_block_idx = worker_group_idx; token_block_idx < num_token_blocks; + token_block_idx += kNumWorkerGroups, ++ iter_idx) + run_task(token_block_idx, iter_idx); + } +}; + +template +struct ScoreBarrier { + static CUTLASS_DEVICE uint64_t encode_state(const uint64_t grid_idx, const uint32_t num_arrived_ctas = 0) { + DG_STATIC_ASSERT(kNumLogicalCtas < layout::mega_gate::kNumMaxLogicalCtas, + "Too many logical CTAs for the score barrier"); + return (grid_idx + 1) * layout::mega_gate::kNumMaxLogicalCtas + num_arrived_ctas; + } + + static CUTLASS_DEVICE void init(uint64_t* state_ptr) { + ptx::st_rel(state_ptr, encode_state(ptx::get_grid_idx())); + } + + static CUTLASS_DEVICE void wait_init(uint64_t* state_ptr) { + const auto state_base = encode_state(ptx::get_grid_idx()); + while (ptx::ld_acq_gpu(state_ptr) - state_base >= kNumLogicalCtas); + } + + static CUTLASS_DEVICE void arrive(uint64_t* state_ptr) { + ptx::red_async_inc_rel(state_ptr); + } + + static CUTLASS_DEVICE void wait(uint64_t* state_ptr) { + const auto all_ctas_arrived_state = encode_state(ptx::get_grid_idx(), kNumLogicalCtas); + while (ptx::ld_acq_gpu(state_ptr) != all_ctas_arrived_state); + } +}; + +} // namespace deep_gemm::sched::mega_gate diff --git a/deep_gemm/include/deep_gemm/scheduler/mega_mhc.cuh b/deep_gemm/include/deep_gemm/scheduler/mega_mhc.cuh new file mode 100644 index 0000000000..8685c446f0 --- /dev/null +++ b/deep_gemm/include/deep_gemm/scheduler/mega_mhc.cuh @@ -0,0 +1,55 @@ +#pragma once + +#include +#include +#include + +namespace deep_gemm::sched::mega_mhc { + +using namespace layout::mega_mhc; + +// Every split arrives once; wait observes all split writes through the +// terminal release sequence. wait_init prevents a producer from racing the grid tag. +template +struct SplitBarrier { + static CUTLASS_DEVICE uint64_t* get_ptr(uint64_t* gmem_split_barriers, const uint32_t m_block_idx) { + const auto line_idx = static_cast(kBarrierIndex) * kNumMaxMBlocks + m_block_idx; + return math::advance_ptr(gmem_split_barriers, line_idx * kSplitBarrierLineBytes); + } + + static CUTLASS_DEVICE uint64_t encode_state(const uint64_t grid_idx, const uint32_t num_arrived_splits = 0) { + return (grid_idx + 1) * kNumSplitBarriers * kNumMaxSplits + num_arrived_splits; + } + + static CUTLASS_DEVICE void init(uint64_t* gmem_split_barriers, const uint32_t& m_block_idx) { + ptx::st_rel(get_ptr(gmem_split_barriers, m_block_idx), encode_state(ptx::get_grid_idx())); + } + + template + static CUTLASS_DEVICE void wait_init(uint64_t* gmem_split_barriers, const uint32_t& m_block_idx) { + const auto current_grid_barrier_base = encode_state(ptx::get_grid_idx()); + const auto split_barrier_ptr = get_ptr(gmem_split_barriers, m_block_idx); + // NOTES: unsigned distance rejects stale grid tags but accepts any current-grid arrival count. + while (ptx::ld_acq_gpu(split_barrier_ptr) - current_grid_barrier_base >= kNumSplits); + } + + template + static CUTLASS_DEVICE void arrive(uint64_t* gmem_split_barriers, const uint32_t& m_block_idx) { + const auto state_ptr = get_ptr(gmem_split_barriers, m_block_idx); + // NOTES: SM100 can retire the release reduction asynchronously while this CTA starts its next task. + ptx::red_async_inc_rel(state_ptr); + } + + template + static CUTLASS_DEVICE void wait(uint64_t* gmem_split_barriers, const uint32_t& m_block_idx) { + const auto all_splits_arrived_state = encode_state(ptx::get_grid_idx(), kNumSplits); + const auto split_barrier_ptr = get_ptr(gmem_split_barriers, m_block_idx); + while (ptx::ld_acq_gpu(split_barrier_ptr) != all_splits_arrived_state); + } +}; + +// Norm is used only by Shifted after materializing X1 and its statistics. +using Norm = SplitBarrier<0>; +using Mix = SplitBarrier<1>; + +} // namespace deep_gemm::sched::mega_mhc diff --git a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh index ba0e11e8f1..c584bcf828 100644 --- a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh @@ -137,6 +137,34 @@ DG_STATIC_ASSERT(sizeof(sched::TaskInfo) == sizeof(sched::TaskInfo) #if defined(__CUDACC__) or defined(__CLION_IDE__) +// Each finished L1 N block toggles its bit in `l2_full_mask`, so L2 K blocks can start as soon as they are fed +// NOTES: the mask parity per ring generation relies on the ring capacity, see `get_num_max_live_pool_blocks` +template +struct L2KBlockDependency { + static constexpr uint32_t kNumL1BlockNs = L1_SHAPE_N / BLOCK_N; + static constexpr uint32_t kNumL1BlockNsPerL2KBlock = BLOCK_K / (BLOCK_N / 2); + static constexpr uint64_t kFullMask = kNumL1BlockNs == 64 ? ~0ull : (1ull << kNumL1BlockNs) - 1; + DG_STATIC_ASSERT(kNumL1BlockNs <= 64, "Too many L1 N blocks for the L2 readiness mask"); + DG_STATIC_ASSERT(BLOCK_K % (BLOCK_N / 2) == 0, "Invalid L1/L2 shape relationship"); + + const uint64_t* mask_ptr; + uint64_t expected_mask, pending_mask = kFullMask; + + CUTLASS_DEVICE L2KBlockDependency(const uint64_t* mask_ptr, const uint32_t& generation_idx): + mask_ptr(mask_ptr), expected_mask(generation_idx & 1 ? 0ull : kFullMask) {} + + CUTLASS_DEVICE static void arrive(const uint64_t* mask_ptr, const uint32_t& n_block_idx) { + ptx::red_xor_rel(mask_ptr, 1ull << n_block_idx); + } + + // Re-read the mask only when the K block is not fed yet + CUTLASS_DEVICE void wait(const uint32_t& k_block_idx) { + const auto k_block_mask = ((1ull << kNumL1BlockNsPerL2KBlock) - 1) << (k_block_idx * kNumL1BlockNsPerL2KBlock); + while (pending_mask & k_block_mask) + pending_mask = ptx::ld_acq_gpu(mask_ptr) ^ expected_mask; + } +}; + template ( num_tokens, lane_idx, workspace.get_shared_l1_task_count_ptr()); + if (is_shared_l2_early) + shared_mainloop( + num_tokens, lane_idx, workspace.get_shared_l2_task_count_ptr()); } // Wait dispatch's results @@ -403,9 +437,9 @@ struct MegaMoEScheduler { } while (task_info.is_valid()); if constexpr (kHasShared) { - // Shared expert L2 tasks depend on SharedLinear1 completion. - shared_mainloop( - num_tokens, lane_idx, workspace.get_shared_l2_task_count_ptr()); + if (not is_shared_l2_early) + shared_mainloop( + num_tokens, lane_idx, workspace.get_shared_l2_task_count_ptr()); } // Sentinel. diff --git a/deep_gemm/include/deep_gemm/scheduler/sm100_mqa_logits.cuh b/deep_gemm/include/deep_gemm/scheduler/sm100_mqa_logits.cuh index 05df973f9c..bafd0b289a 100644 --- a/deep_gemm/include/deep_gemm/scheduler/sm100_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/sm100_mqa_logits.cuh @@ -4,14 +4,51 @@ #include // SM100 contiguous-KV scheduler; translates row bounds into core task geometry +// Two scheduler modes: +// - Grid-stride: SM `s` owns Q blocks `s`, `s + kNumSMs`, never split K +// - With metadata: consumes metadata prebuilt by sm100_mqa_logits_metadata, +// may split K to equalize workload among SMs. Metadata workspace layout, count as int32: +// [0, 2 * kNumSMs) per-SM start points, uint2 = (q_block_idx, split_offset) +// [2 * kNumSMs, 3 * kNumSMs) per-SM split counts +// [3 * kNumSMs, kMetaSpanWordOffset) pad kMetaSpanWordOffset to 8-byte aligned +// [kMetaSpanWordOffset, ...) per-block MQALogitsKVSpan namespace deep_gemm::sched { -template +// KV range one Q block must visit, cover all [ks, ke) +// precomputed by metadata kernel, consumed by (Q TMA, KV TMA, UMMA, math) +struct MQALogitsKVSpan { + uint32_t kv_token_base; // min k_start over rows, rounded down to 4 tokens + uint32_t num_kv_splits; // ceil((max k_end - kv_token_base) / SPLIT_KV) +}; + +// NOTE: Out-of-range rows are clamped to the last valid row +template +CUTLASS_DEVICE MQALogitsKVSpan get_mqa_logits_kv_span(const uint32_t& q_block_idx, + const uint32_t& num_q_tokens, const uint32_t& num_kv_tokens, + const uint32_t* cu_seq_len_k_start, + const uint32_t* cu_seq_len_k_end) { + uint32_t start = cute::numeric_limits::max(); + uint32_t end = cute::numeric_limits::min(); + #pragma unroll 8 + for (uint32_t token_idx = 0; token_idx < BLOCK_Q; ++ token_idx) { + const auto row_idx = cute::min(q_block_idx * BLOCK_Q + token_idx, num_q_tokens - 1); + const auto k_start = cute::min(cu_seq_len_k_start[row_idx], num_kv_tokens); + const auto k_end = cute::min(cu_seq_len_k_end[row_idx], num_kv_tokens); + start = cute::min(start, k_start); + end = cute::max(end, k_end); + } + const uint32_t kv_token_base = start / 4 * 4; + return {kv_token_base, math::ceil_div(end - kv_token_base, SPLIT_KV)}; +} + +template struct SM100MQALogitsScheduler { static constexpr bool kIsPaged = false; static constexpr bool kHasPartialBlock = false; static constexpr uint32_t kPageKV = 0; + static constexpr uint32_t kMetaSpanWordOffset = (3 * kNumSMs + 1) / 2 * 2; uint32_t current_q_block_idx; uint32_t num_q_blocks; @@ -19,47 +56,92 @@ struct SM100MQALogitsScheduler { uint32_t num_kv_tokens; const uint32_t* cu_seq_len_k_start; const uint32_t* cu_seq_len_k_end; - uint32_t* seq_k_start; - uint32_t* seq_k_end; + uint32_t current_split_offset = 0; + uint32_t remaining_splits = 0; + const MQALogitsKVSpan* kv_spans = nullptr; CUTLASS_DEVICE SM100MQALogitsScheduler(const uint32_t& sm_idx, const uint32_t& num_q_tokens, const uint32_t& num_kv_tokens, const uint32_t* cu_seq_len_k_start, const uint32_t* cu_seq_len_k_end, - uint32_t* seq_k_start, - uint32_t* seq_k_end): + const uint32_t* schedule_meta = nullptr): current_q_block_idx(sm_idx), num_q_blocks(math::ceil_div(num_q_tokens, BLOCK_Q)), num_q_tokens(num_q_tokens), num_kv_tokens(num_kv_tokens), cu_seq_len_k_start(cu_seq_len_k_start), - cu_seq_len_k_end(cu_seq_len_k_end), - seq_k_start(seq_k_start), - seq_k_end(seq_k_end) {} - - // Contiguous KV uses absolute token offsets; align the base to 4 for compressed logits - CUTLASS_DEVICE bool next_q_block(uint32_t& q_block_idx, uint32_t& kv_token_base, uint32_t& num_kv_splits) { - if (current_q_block_idx >= num_q_blocks) - return false; + cu_seq_len_k_end(cu_seq_len_k_end) { + DG_STATIC_ASSERT(kNumSMs > 0, "Invalid SM count"); + if constexpr (kUseSchedule) { + const auto start = reinterpret_cast(schedule_meta)[sm_idx]; + current_q_block_idx = start.x; + current_split_offset = start.y; + remaining_splits = schedule_meta[2 * kNumSMs + sm_idx]; + kv_spans = reinterpret_cast(schedule_meta + kMetaSpanWordOffset); + } + } - q_block_idx = current_q_block_idx; - current_q_block_idx += kNumSMs; + CUTLASS_DEVICE auto make_cleaner(const uint32_t& sm_idx) const { + return SM100MQALogitsScheduler( + sm_idx, num_q_tokens, num_kv_tokens, cu_seq_len_k_start, cu_seq_len_k_end); + } - uint32_t start = cute::numeric_limits::max(); - uint32_t end = cute::numeric_limits::min(); - #pragma unroll - for (uint32_t token_idx = 0; token_idx < BLOCK_Q; ++ token_idx) { - const auto row_idx = cute::min(q_block_idx * BLOCK_Q + token_idx, num_q_tokens - 1); - seq_k_start[token_idx] = cute::min(cu_seq_len_k_start[row_idx], num_kv_tokens); - seq_k_end[token_idx] = cute::min(cu_seq_len_k_end[row_idx], num_kv_tokens); - start = cute::min(start, seq_k_start[token_idx]); - end = cute::max(end, seq_k_end[token_idx]); + template + CUTLASS_DEVICE bool next_q_block(uint32_t& q_block_idx, uint32_t& kv_token_base, uint32_t& num_kv_splits, + uint32_t* seq_k_start = nullptr, uint32_t* seq_k_end = nullptr) { + if constexpr (kUseSchedule) { + while (remaining_splits > 0 and current_q_block_idx < num_q_blocks) { + const auto span = kv_spans[current_q_block_idx]; + const auto split_offset = current_split_offset; + current_split_offset = 0; + if (split_offset < span.num_kv_splits) { + if constexpr (kLoadSeqBounds) { + #pragma unroll + for (uint32_t token_idx = 0; token_idx < BLOCK_Q; ++ token_idx) { + const auto row_idx = cute::min(current_q_block_idx * BLOCK_Q + token_idx, num_q_tokens - 1); + seq_k_start[token_idx] = cute::min(cu_seq_len_k_start[row_idx], num_kv_tokens); + seq_k_end[token_idx] = cute::min(cu_seq_len_k_end[row_idx], num_kv_tokens); + } + } + q_block_idx = current_q_block_idx; + kv_token_base = span.kv_token_base + split_offset * SPLIT_KV; + num_kv_splits = cute::min(span.num_kv_splits - split_offset, remaining_splits); + remaining_splits -= num_kv_splits; + ++ current_q_block_idx; + return true; + } + ++ current_q_block_idx; + } + return false; + } else { + if (current_q_block_idx >= num_q_blocks) + return false; + q_block_idx = current_q_block_idx; + current_q_block_idx += kNumSMs; + if constexpr (kLoadSeqBounds) { + uint32_t start = cute::numeric_limits::max(); + uint32_t end = cute::numeric_limits::min(); + #pragma unroll + for (uint32_t token_idx = 0; token_idx < BLOCK_Q; ++ token_idx) { + const auto row_idx = cute::min(q_block_idx * BLOCK_Q + token_idx, num_q_tokens - 1); + const auto k_start = cute::min(cu_seq_len_k_start[row_idx], num_kv_tokens); + const auto k_end = cute::min(cu_seq_len_k_end[row_idx], num_kv_tokens); + seq_k_start[token_idx] = k_start; + seq_k_end[token_idx] = k_end; + start = cute::min(start, k_start); + end = cute::max(end, k_end); + } + kv_token_base = start / 4 * 4; + num_kv_splits = math::ceil_div(end - kv_token_base, SPLIT_KV); + } else { + const auto span = get_mqa_logits_kv_span(q_block_idx, num_q_tokens, num_kv_tokens, + cu_seq_len_k_start, cu_seq_len_k_end); + kv_token_base = span.kv_token_base; + num_kv_splits = span.num_kv_splits; + } + return true; } - - kv_token_base = start / 4 * 4; - num_kv_splits = math::ceil_div(end - kv_token_base, SPLIT_KV); - return true; } CUTLASS_DEVICE uint32_t get_q_tma_token_base(const uint32_t& q_block_idx) const { diff --git a/deep_gemm/include/deep_gemm/scheduler/sm100_mqa_logits_metadata.cuh b/deep_gemm/include/deep_gemm/scheduler/sm100_mqa_logits_metadata.cuh new file mode 100644 index 0000000000..a3b7c6ff18 --- /dev/null +++ b/deep_gemm/include/deep_gemm/scheduler/sm100_mqa_logits_metadata.cuh @@ -0,0 +1,121 @@ +#pragma once + +#include + +#include +#include +#include +#include + +// SM100 scheduler for contiguous-KV logits emits per-SM start point (q_block, split_offset) and work counts +// Each Q block priced as kMQALogitsSplitCost * num_kv_splits + kMQALogitsSegmentCost +// both constants are 1 for now + +namespace deep_gemm::sched { + +inline constexpr uint32_t kNumLogitsMetadataThreads = 256; +inline constexpr uint32_t kNumLogitsMetadataWarps = kNumLogitsMetadataThreads / 32; +DG_STATIC_ASSERT(kNumLogitsMetadataWarps <= 32 and kNumLogitsMetadataThreads % 32 == 0, "Invalid metadata thread count"); + +// Number of prefix <= target, at most count +template +CUTLASS_DEVICE uint32_t mqa_logits_prefix_upper_bound(const uint32_t* prefix, const uint32_t& count, + const uint32_t& target) { + uint32_t pos = 0; + #pragma unroll + for (uint32_t step = 1u << (kMaxLog2 - 1); step > 0; step >>= 1) { + const uint32_t candidate = pos + step; + if (candidate <= count and prefix[candidate - 1] <= target) + pos = candidate; + } + return pos; +} + +template +CUTLASS_GLOBAL __launch_bounds__(kNumLogitsMetadataThreads, 1) +void sm100_mqa_logits_metadata(const uint32_t num_q_tokens, + const uint32_t num_kv_tokens, + const uint32_t* cu_seq_len_k_start, + const uint32_t* cu_seq_len_k_end, + uint32_t* schedule_meta) { + DG_STATIC_ASSERT(kSegmentCost > 0 and kSplitCost > 0, "Invalid costs"); + DG_STATIC_ASSERT(kNumSMs <= kNumLogitsMetadataThreads, "Invalid SM count"); + constexpr uint32_t kSpanWordOffset = SM100MQALogitsScheduler::kMetaSpanWordOffset; + + const uint32_t thread_idx = threadIdx.x; + DG_DEVICE_ASSERT(blockDim.x == kNumLogitsMetadataThreads); + cudaGridDependencySynchronize(); + if (cutlass::canonical_warp_idx_sync() == 0 and cute::elect_one_sync()) + cutlass::arch::launch_dependent_grids(); + + const uint32_t num_q_blocks = math::ceil_div(num_q_tokens, BLOCK_Q); + extern __shared__ uint32_t metadata_storage[]; + const auto work_prefix = metadata_storage; // [num_q_blocks] + const auto cost_prefix = work_prefix + num_q_blocks; // [num_q_blocks] + const auto kv_spans = reinterpret_cast(schedule_meta + kSpanWordOffset); // [num_q_blocks] + + for (uint32_t q_block_idx = thread_idx; q_block_idx < num_q_blocks; q_block_idx += kNumLogitsMetadataThreads) { + const auto span = get_mqa_logits_kv_span(q_block_idx, num_q_tokens, num_kv_tokens, + cu_seq_len_k_start, cu_seq_len_k_end); + kv_spans[q_block_idx] = span; + work_prefix[q_block_idx] = span.num_kv_splits; + cost_prefix[q_block_idx] = kSplitCost * span.num_kv_splits + (span.num_kv_splits == 0 ? 0 : kSegmentCost); + } + __syncthreads(); + + __shared__ uint64_t warp_sums[kNumLogitsMetadataWarps]; + __shared__ uint64_t carry; + if (thread_idx == 0) + carry = 0; + for (uint32_t base_q_block_idx = 0; base_q_block_idx < num_q_blocks; base_q_block_idx += kNumLogitsMetadataThreads) { + const uint32_t q_block_idx = base_q_block_idx + thread_idx; + const bool active = q_block_idx < num_q_blocks; + // One scan for both prefixes, work in high half, cost in the low + // both bounded to u32, will not interfere + const uint64_t value = active ? (static_cast(work_prefix[q_block_idx]) << 32) + cost_prefix[q_block_idx] : 0; + // publish the carry, protect warp_sums + __syncthreads(); + // no data race: read before the scan call, whose internal __syncthreads() separate this read from the next write + const uint64_t previous_carry = carry; + const uint64_t scanned = math::cta_exclusive_sum(value, warp_sums) + value + previous_carry; + if (active) { + work_prefix[q_block_idx] = static_cast(scanned >> 32); + cost_prefix[q_block_idx] = static_cast(scanned); + } + if (thread_idx == kNumLogitsMetadataThreads - 1) + carry = scanned; + __syncthreads(); + } + // num_q_blocks > 0 is asserted on the host side + const uint32_t total_work = work_prefix[num_q_blocks - 1]; + const uint32_t total_cost = cost_prefix[num_q_blocks - 1]; + + if (thread_idx >= kNumSMs) + return; + const uint32_t base = total_cost / kNumSMs; + const uint32_t remainder = total_cost % kNumSMs; + struct Boundary { + uint2 start; + uint32_t coordinate; + }; + const auto locate = [&](const uint32_t& boundary_idx) -> Boundary { + const uint32_t target = boundary_idx * base + cute::min(boundary_idx, remainder); + if (target == total_cost) + return {make_uint2(num_q_blocks, 0), total_work}; // Tail sentinel: one-past-the-end + const uint32_t q_block_idx = mqa_logits_prefix_upper_bound(cost_prefix, num_q_blocks, target); + const uint32_t cost_before = q_block_idx == 0 ? 0 : cost_prefix[q_block_idx - 1]; + const uint32_t work_before = q_block_idx == 0 ? 0 : work_prefix[q_block_idx - 1]; + // Targets inside a block land after its segment charge; the clamp + // keeps the boundary before the block's last split + const uint32_t split_offset = cute::min((cute::max(target - cost_before, kSegmentCost) - kSegmentCost) / kSplitCost, + work_prefix[q_block_idx] - work_before - 1); + return {make_uint2(q_block_idx, split_offset), work_before + split_offset}; + }; + const auto begin = locate(thread_idx); + const auto end = locate(thread_idx + 1); + reinterpret_cast(schedule_meta)[thread_idx] = begin.start; + schedule_meta[2 * kNumSMs + thread_idx] = end.coordinate - begin.coordinate; +} + +} // namespace deep_gemm::sched diff --git a/deep_gemm/include/deep_gemm/scheduler/sm100_paged_mqa_logits.cuh b/deep_gemm/include/deep_gemm/scheduler/sm100_paged_mqa_logits.cuh index 3c6e9a6cb0..8306c78861 100644 --- a/deep_gemm/include/deep_gemm/scheduler/sm100_paged_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/sm100_paged_mqa_logits.cuh @@ -19,33 +19,32 @@ struct RequestInfo { uint32_t num_kv_splits; // request_num_kv_splits = ceil(context_len / SPLIT_KV) uint32_t num_kv_pages; // = ceil(context_len / PAGE_KV); page-level bound for the last partial split + CUTLASS_DEVICE RequestInfo() = default; + + CUTLASS_DEVICE RequestInfo(const uint32_t& q_token_start, const uint32_t& num_q_tokens, + const uint32_t& context_len) + : q_token_start(q_token_start), num_q_tokens(num_q_tokens), + num_q_blocks(math::ceil_div(num_q_tokens, BLOCK_Q)), + num_kv_splits(math::ceil_div(context_len, SPLIT_KV)), + num_kv_pages(math::ceil_div(context_len, PAGE_KV)) {} + // Resolve the request that starts at `q_token_idx` CUTLASS_DEVICE static RequestInfo from_q_token(const uint32_t& q_token_idx, const uint32_t& num_q_tokens_total, const uint32_t* context_lens, const uint32_t* indices) { - RequestInfo info; - info.q_token_start = q_token_idx; - uint32_t context_len; if constexpr (kIsVarlen) { // Varlen request = maximal run of equal `indices` const uint32_t request_id = indices[q_token_idx]; uint32_t t = q_token_idx; while (t + 1 < num_q_tokens_total and indices[t + 1] == request_id) ++ t; - info.num_q_tokens = t - q_token_idx + 1; - context_len = context_lens[t]; + return RequestInfo(q_token_idx, t - q_token_idx + 1, context_lens[t]); } else { // Regular grid: request = q_token_idx / next_n, next_n tokens each - const uint32_t request_id = q_token_idx / kNextN; - info.num_q_tokens = kNextN; - const uint32_t lens_idx = kIsContextLens2D ? request_id * kNextN + kNextN - 1 : request_id; - context_len = context_lens[lens_idx]; + const uint32_t lens_idx = kIsContextLens2D ? q_token_idx + kNextN - 1 : q_token_idx / kNextN; + return RequestInfo(q_token_idx, kNextN, context_lens[lens_idx]); } - info.num_q_blocks = math::ceil_div(info.num_q_tokens, BLOCK_Q); - info.num_kv_splits = math::ceil_div(context_len, SPLIT_KV); - info.num_kv_pages = math::ceil_div(context_len, PAGE_KV); - return info; } // Average q-token partition across Q-blocks; returns both offset and count @@ -64,99 +63,155 @@ struct RequestInfo { } }; -// Metadata kernel balances work across SMs via a prefix sum over request work +inline constexpr uint32_t kNumMetadataThreads = 1024; +inline constexpr uint32_t kNumMetadataWarps = kNumMetadataThreads / 32; +DG_STATIC_ASSERT(kNumMetadataWarps <= 32 and kNumMetadataThreads % 32 == 0, "Invalid metadata thread count"); + +// Inclusive prefix sum +CUTLASS_DEVICE void metadata_prefix_scan(const uint32_t thread_idx, const uint32_t num_items, + uint32_t* values, uint32_t* warp_sums) { + const uint32_t num_items_per_thread = math::ceil_div(num_items, kNumMetadataThreads) | 1u; + const uint32_t item_begin_idx = cute::min(thread_idx * num_items_per_thread, num_items); + const uint32_t item_end_idx = cute::min(item_begin_idx + num_items_per_thread, num_items); + uint32_t even_sum = 0, odd_sum = 0; + uint32_t item_idx = item_begin_idx; + for (; item_idx + 2 <= item_end_idx; item_idx += 2) { + even_sum += values[item_idx]; + values[item_idx] = even_sum + odd_sum; + odd_sum += values[item_idx + 1]; + values[item_idx + 1] = odd_sum + even_sum; + } + if (item_idx < item_end_idx) { + even_sum += values[item_idx]; + values[item_idx] = even_sum + odd_sum; + } + const uint32_t thread_offset = math::cta_exclusive_sum(even_sum + odd_sum, warp_sums); + for (item_idx = item_begin_idx; item_idx < item_end_idx; ++item_idx) + values[item_idx] += thread_offset; + __syncthreads(); +} + +// Balance work across SMs by request prefix sum template -CUTLASS_GLOBAL __launch_bounds__(256, 1) + uint32_t SPLIT_KV, uint32_t kNumSMs> +CUTLASS_GLOBAL __launch_bounds__(kNumMetadataThreads, 1) void sm100_paged_mqa_logits_metadata(const uint32_t num_requests, const uint32_t num_q_tokens_total, const uint32_t* context_lens, const uint32_t* indices, uint32_t* schedule_meta) { - // PAGE_KV is unused for metadata; pass SPLIT_KV as a placeholder - using Info = RequestInfo; - - cudaGridDependencySynchronize(); // wait for the primary kernel (CDP launch) - const uint32_t thread_idx = threadIdx.x; - const uint32_t lane_idx = ptx::get_lane_idx(); - const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); - const uint32_t num_threads = blockDim.x; + DG_DEVICE_ASSERT(blockDim.x == kNumMetadataThreads); + cudaGridDependencySynchronize(); // wait for the primary kernel (CDP launch) - // smem: per-request work prefix sum + request start token extern __shared__ uint32_t smem[]; - uint32_t* prefix_work = smem; // [num_requests] - uint32_t* request_q_token_start = smem + num_requests; // [num_requests] + const auto request_q_token_start_idx = smem; // [num_requests], varlen only + const auto request_work_prefix = smem + (kIsVarlen ? num_requests : 0); // [num_requests] + const auto warp_sums = request_work_prefix + num_requests; // [kNumMetadataWarps] + const auto num_logical_requests_shared = warp_sums + kNumMetadataWarps; - // Logical requests are regular requests for non-varlen and index runs for varlen uint32_t num_logical_requests; if constexpr (kIsVarlen) { - if (thread_idx == 0) { - uint32_t r = 0, t = 0; - while (t < num_q_tokens_total) { - request_q_token_start[r] = t; - const uint32_t request_id = indices[t]; - while (t < num_q_tokens_total and indices[t] == request_id) - ++ t; - ++ r; - } - // Temporarily stash run count for broadcast - prefix_work[0] = r; // temporary: run count + // Extract request starts + DG_DEVICE_ASSERT(reinterpret_cast(indices) % 8 == 0); + const uint32_t num_tokens_per_thread = (math::ceil_div(num_q_tokens_total, kNumMetadataThreads * 2) | 1u) * 2; + const uint32_t lo = cute::min(thread_idx * num_tokens_per_thread, num_q_tokens_total); + const uint32_t hi = cute::min(lo + num_tokens_per_thread, num_q_tokens_total); + uint32_t num_request_starts = 0; + uint32_t prev_id = lo > 0 and lo < hi ? indices[lo - 1] : 0u; + const auto indices_vec2 = reinterpret_cast(indices); + uint32_t token_idx = lo; + for (; token_idx + 2 <= hi; token_idx += 2) { + const uint2 ids = indices_vec2[token_idx / 2]; + const bool is_x_start = (token_idx == 0) or (ids.x != prev_id); + const bool is_y_start = ids.y != ids.x; + num_request_starts += is_x_start + is_y_start; + prev_id = ids.y; } + if (token_idx < hi) { + const bool is_request_start = (token_idx == 0) or (indices[token_idx] != prev_id); + num_request_starts += is_request_start; + } + uint32_t request_idx = math::cta_exclusive_sum(num_request_starts, warp_sums); + prev_id = lo > 0 and lo < hi ? indices[lo - 1] : 0u; + token_idx = lo; + for (; token_idx + 2 <= hi; token_idx += 2) { + const uint2 ids = indices_vec2[token_idx / 2]; + const bool is_x_start = (token_idx == 0) or (ids.x != prev_id); + const bool is_y_start = ids.y != ids.x; + if (is_x_start) + request_q_token_start_idx[request_idx ++] = token_idx; + if (is_y_start) + request_q_token_start_idx[request_idx ++] = token_idx + 1; + prev_id = ids.y; + } + if (token_idx < hi) { + const bool is_request_start = (token_idx == 0) or (indices[token_idx] != prev_id); + if (is_request_start) + request_q_token_start_idx[request_idx ++] = token_idx; + } + if (thread_idx == kNumMetadataThreads - 1) + *num_logical_requests_shared = request_idx; __syncthreads(); - num_logical_requests = prefix_work[0]; - __syncthreads(); + num_logical_requests = *num_logical_requests_shared; } else { num_logical_requests = num_requests; - for (uint32_t r = thread_idx; r < num_logical_requests; r += num_threads) - request_q_token_start[r] = r * kNextN; - __syncthreads(); } - // Work per request before prefix sum - for (uint32_t r = thread_idx; r < num_logical_requests; r += num_threads) { - const auto info = Info::from_q_token(request_q_token_start[r], - num_q_tokens_total, context_lens, indices); - prefix_work[r] = info.num_kv_splits * info.num_q_tokens; - } - __syncthreads(); - - // Inclusive prefix sum by one warp - if (warp_idx == 0) { - uint32_t carry = 0; - for (uint32_t base = 0; base < num_logical_requests; base += 32) { - const uint32_t r = base + lane_idx; - const uint32_t v = (r < num_logical_requests) ? prefix_work[r] : 0u; - const uint32_t scanned = math::warp_inclusive_sum(v, lane_idx) + carry; - if (r < num_logical_requests) - prefix_work[r] = scanned; - carry = __shfl_sync(0xffffffff, scanned, 31); + const auto get_request_info = [&](const uint32_t& request_idx, + uint32_t& q_token_start_idx, + uint32_t& num_q_tokens, + uint32_t& context_len) { + if constexpr (kIsVarlen) { + q_token_start_idx = request_q_token_start_idx[request_idx]; + const uint32_t q_token_end_idx = request_idx + 1 < num_logical_requests ? + request_q_token_start_idx[request_idx + 1] : num_q_tokens_total; + num_q_tokens = q_token_end_idx - q_token_start_idx; + context_len = context_lens[q_token_end_idx - 1]; + } else { + q_token_start_idx = request_idx * kNextN; + num_q_tokens = kNextN; + const uint32_t lens_idx = kIsContextLens2D ? request_idx * kNextN + kNextN - 1 : request_idx; + context_len = context_lens[lens_idx]; } + }; + + // Compute per-request work + for (uint32_t request_idx = thread_idx; request_idx < num_logical_requests; request_idx += kNumMetadataThreads) { + uint32_t q_token_start_idx, num_q_tokens, context_len; + get_request_info(request_idx, q_token_start_idx, num_q_tokens, context_len); + request_work_prefix[request_idx] = math::ceil_div(context_len, SPLIT_KV) * num_q_tokens; } __syncthreads(); - const uint32_t num_total_work = num_logical_requests > 0 ? prefix_work[num_logical_requests - 1] : 0u; - - // Each thread emits one SM start; remainder is assigned to earlier SMs - const uint32_t q = num_total_work / kNumSMs, rem = num_total_work % kNumSMs; - for (uint32_t sm_idx = thread_idx; sm_idx <= kNumSMs; sm_idx += num_threads) { - const uint32_t w = sm_idx * q + (sm_idx < rem ? sm_idx : rem); - // First request whose prefix_work owns work unit `w` + if (num_logical_requests > 0) { + metadata_prefix_scan(thread_idx, num_logical_requests, request_work_prefix, warp_sums); + } + const uint32_t num_total_work = num_logical_requests > 0 ? request_work_prefix[num_logical_requests - 1] : 0u; + + // Partition work across SMs + const uint32_t q = num_total_work / kNumSMs; + const uint32_t rem = num_total_work % kNumSMs; + for (uint32_t sm_idx = thread_idx; sm_idx <= kNumSMs; sm_idx += kNumMetadataThreads) { + const uint32_t w = sm_idx * q + cute::min(sm_idx, rem); + // Find the request containing `w` uint32_t lo = 0, hi = num_logical_requests; while (lo < hi) { const uint32_t mid = (lo + hi) / 2; - if (prefix_work[mid] <= w) lo = mid + 1; else hi = mid; + if (request_work_prefix[mid] <= w) + lo = mid + 1; + else + hi = mid; } const uint32_t request_idx = lo; uint32_t q_token_idx, kv_split_idx; if (request_idx < num_logical_requests) { - const uint32_t work_before = (request_idx == 0) ? 0u : prefix_work[request_idx - 1]; - const uint32_t w_in_request = w - work_before; - const auto info = Info::from_q_token(request_q_token_start[request_idx], - num_q_tokens_total, context_lens, indices); + const uint32_t w_in_request = w - (request_idx == 0 ? 0u : request_work_prefix[request_idx - 1]); + uint32_t q_token_start_idx, num_q_tokens, context_len; + get_request_info(request_idx, q_token_start_idx, num_q_tokens, context_len); // Align SM starts to request/split boundaries - q_token_idx = info.q_token_start; - kv_split_idx = w_in_request / info.num_q_tokens; + q_token_idx = q_token_start_idx; + kv_split_idx = w_in_request / num_q_tokens; } else { // Tail sentinel: one-past-the-end q_token_idx = num_q_tokens_total; @@ -246,7 +301,9 @@ struct SM100PagedMQALogitsScheduler : SM100IndicesStorage { } // Emit the next (Q-block, chunk) task and stash its addressing geometry - CUTLASS_DEVICE bool next_q_block(uint32_t& q_block_idx, uint32_t& kv_split_base, uint32_t& num_kv_splits) { + template + CUTLASS_DEVICE bool next_q_block(uint32_t& q_block_idx, uint32_t& kv_split_base, uint32_t& num_kv_splits, + uint32_t* = nullptr, uint32_t* = nullptr) { q_block_idx = 0; // addressing uses stashed state if (done) return false; diff --git a/deep_gemm/include/deep_gemm/scheduler/sm100_sparse_mqa_logits_metadata.cuh b/deep_gemm/include/deep_gemm/scheduler/sm100_sparse_mqa_logits_metadata.cuh new file mode 100644 index 0000000000..ec5dbd1884 --- /dev/null +++ b/deep_gemm/include/deep_gemm/scheduler/sm100_sparse_mqa_logits_metadata.cuh @@ -0,0 +1,468 @@ +#pragma once + +#include + +#include +#include +#include +#include + +namespace deep_gemm::sched::sparse_mqa_logits { + +using namespace layout::sparse_mqa_logits; + +template +struct SharedStorage { + static constexpr uint32_t kNumMaxMergedKVBlocks = BLOCK_Q * NUM_MAX_SPARSE_BLOCKS; + static constexpr uint32_t kNumMaxKVSplits = math::constexpr_ceil_div(kNumMaxMergedKVBlocks, kNumKVBlocksPerSplit); + static constexpr uint32_t kNumWarps = kNumThreads / 32; + + uint32_t logical_kv_block_indices[BLOCK_Q][NUM_MAX_SPARSE_BLOCKS]; + uint32_t packed_slots_by_merged_kv_block[kNumMaxMergedKVBlocks]; + uint32_t packed_slot_bases_by_kv_split[kNumMaxKVSplits]; + uint32_t warp_sums[kNumWarps]; + uint32_t num_waves; + uint32_t kv_split_histograms[kNumWarps][kNumKVSplitsPerEntry + 1]; + struct alignas(16) { + uint32_t q_token_base; + uint32_t num_q_tokens; + uint32_t num_kv_blocks[BLOCK_Q]; + uint32_t kv_split_base; + } q_block; + uint32_t is_last_cta; +}; + +// Divide KV splits evenly, then split each SM range at Q-block boundaries. +template +CUTLASS_DEVICE uint32_t build_contiguous_schedule( + ScheduleEntry* schedule_entries, const KVSplit* kv_splits, + const QBlockInfo* q_block_infos, const uint32_t num_q_tokens, + const uint32_t total_kv_splits, smem_t& smem +) { + if (threadIdx.x == 0) + smem.num_waves = 1; + __syncthreads(); + + uint32_t num_sm_entries = 0; + if (threadIdx.x < kNumSMs) { + uint32_t kv_split_idx = static_cast(math::ceil_div( + static_cast(total_kv_splits) * threadIdx.x, kNumSMs)); + const uint32_t kv_split_end = static_cast(math::ceil_div( + static_cast(total_kv_splits) * (threadIdx.x + 1), kNumSMs)); + while (kv_split_idx < kv_split_end) { + const uint32_t q_token_base = kv_splits[kv_split_idx].header.q_token_base; + const auto q_block_info = q_block_infos[q_token_base]; + const uint32_t entry_kv_split_end = cute::min(kv_split_end, q_block_info.kv_split_base + q_block_info.num_kv_splits); + schedule_entries[num_sm_entries * kNumSMs + threadIdx.x] = ScheduleEntry( + kv_split_idx, entry_kv_split_end, q_token_base, cute::min(BLOCK_Q, num_q_tokens - q_token_base)); + kv_split_idx = entry_kv_split_end; + ++ num_sm_entries; + } + atomicMax(&smem.num_waves, num_sm_entries); + } + __syncthreads(); + + const uint32_t num_waves = smem.num_waves; + if (threadIdx.x < kNumSMs) { + for (uint32_t wave_idx = num_sm_entries; wave_idx < num_waves; ++ wave_idx) + schedule_entries[wave_idx * kNumSMs + threadIdx.x] = ScheduleEntry(0, 0, 0, 0); + } + return num_waves; +} + +// Sort each wave by KV-split count and rotate heavy entries across SMs. +template +CUTLASS_DEVICE void balance_wave_entries(ScheduleEntry* schedule_entries, const uint32_t num_waves, smem_t& smem) { + __syncthreads(); + if (num_waves == 1) + return; + const uint32_t lane_idx = ptx::get_lane_idx(); + const uint32_t warp_idx = cutlass::canonical_warp_idx_sync(); + + for (uint32_t wave_idx = warp_idx; wave_idx < num_waves; wave_idx += smem_t::kNumWarps) { + constexpr uint32_t kNumEntriesPerLane = math::constexpr_ceil_div(kNumSMs, 32u); + const auto histogram = smem.kv_split_histograms[warp_idx]; + const auto wave_entries = schedule_entries + wave_idx * kNumSMs; + ScheduleEntry entries[kNumEntriesPerLane]; + uint32_t ranks[kNumEntriesPerLane]; + + #pragma unroll + for (uint32_t sm_idx = lane_idx, entry_idx = 0; sm_idx < kNumSMs; sm_idx += 32, ++ entry_idx) + entries[entry_idx] = wave_entries[sm_idx]; + for (uint32_t num_kv_splits = lane_idx; num_kv_splits <= kNumKVSplitsPerEntry; num_kv_splits += 32) + histogram[num_kv_splits] = 0; + __syncwarp(); + + #pragma unroll + for (uint32_t sm_idx = lane_idx, entry_idx = 0; sm_idx < kNumSMs; sm_idx += 32, ++ entry_idx) + ranks[entry_idx] = atomicAdd(histogram + entries[entry_idx].kv_split_end - entries[entry_idx].kv_split_begin, 1u); + __syncwarp(); + + if (cute::elect_one_sync()) { + uint32_t rank_begin = 0; + for (uint32_t num_kv_splits = 0; num_kv_splits <= kNumKVSplitsPerEntry; ++ num_kv_splits) { + const uint32_t frequency = histogram[num_kv_splits]; + histogram[num_kv_splits] = rank_begin; + rank_begin += frequency; + } + } + __syncwarp(); + + #pragma unroll + for (uint32_t sm_idx = lane_idx, entry_idx = 0; sm_idx < kNumSMs; sm_idx += 32, ++ entry_idx) { + const uint32_t rank = histogram[entries[entry_idx].kv_split_end - entries[entry_idx].kv_split_begin] + ranks[entry_idx]; + // Reverse the two-wave tail to make active SM sets complementary. + const uint32_t dst_sm_idx = num_waves == 2 and wave_idx == 1 ? kNumSMs - 1 - rank : + (rank + kNumSMs - wave_idx * kNumSMs / num_waves) % kNumSMs; + wave_entries[dst_sm_idx] = entries[entry_idx]; + } + } +} + +// Split each Q block into bounded entries, then balance each wave. +template +CUTLASS_DEVICE uint32_t build_paged_schedule( + ScheduleEntry* schedule_entries, const QBlockInfo* q_block_infos, + const uint32_t num_q_tokens, const uint32_t* indices, smem_t& smem +) { + uint32_t num_entries = 0; + for (uint32_t q_token_begin = 0; q_token_begin < num_q_tokens; q_token_begin += kNumThreads) { + const uint32_t q_token_idx = q_token_begin + threadIdx.x; + const auto q_block_info = q_token_idx < num_q_tokens ? q_block_infos[q_token_idx] : QBlockInfo(0, 0); + const uint32_t num_kv_splits = q_block_info.num_kv_splits; + const uint32_t num_q_entries = math::ceil_div(num_kv_splits, kNumKVSplitsPerEntry); + uint32_t num_batch_entries; + const uint32_t entry_begin = num_entries + math::cta_exclusive_sum( + num_q_entries, smem.warp_sums, num_batch_entries); + + const uint32_t kv_splits_per_entry = num_q_entries == 0 ? 0 : num_kv_splits / num_q_entries; + const uint32_t num_larger_entries = num_q_entries == 0 ? 0 : num_kv_splits % num_q_entries; + const uint32_t num_q_block_tokens = q_token_idx + 1 < num_q_tokens and indices[q_token_idx + 1] == indices[q_token_idx] ? BLOCK_Q : 1; + uint32_t kv_split_begin = q_block_info.kv_split_base; + for (uint32_t q_entry_idx = 0; q_entry_idx < num_q_entries; ++ q_entry_idx) { + const uint32_t kv_split_end = kv_split_begin + kv_splits_per_entry + (q_entry_idx < num_larger_entries); + schedule_entries[entry_begin + q_entry_idx] = ScheduleEntry( + kv_split_begin, kv_split_end, q_token_idx, num_q_block_tokens); + kv_split_begin = kv_split_end; + } + num_entries += num_batch_entries; + __syncthreads(); + } + + const uint32_t num_waves = cute::max(1u, math::ceil_div(num_entries, kNumSMs)); + for (uint32_t entry_idx = num_entries + threadIdx.x; entry_idx < num_waves * kNumSMs; entry_idx += kNumThreads) + schedule_entries[entry_idx] = ScheduleEntry(0, 0, 0, 0); + + balance_wave_entries(schedule_entries, num_waves, smem); + return num_waves; +} + +template +CUTLASS_GLOBAL __launch_bounds__(kNumThreads, 4) +void sm100_sparse_mqa_logits_metadata( + const uint32_t num_q_tokens, const uint32_t num_kv_tokens, + const uint32_t* cu_seq_len_k_start, const uint32_t* cu_seq_len_k_end, + const uint32_t* context_lens, const uint32_t* block_table, const uint32_t block_table_stride, + const uint32_t* indices, + const uint32_t* sparse_kv_block_indices, uint8_t* metadata, uint8_t* workspace +) { + // Shared memory sizes + constexpr uint32_t kNumKVBlocksPerSplit = SPLIT_KV / SPARSE_BLOCK_KV; + using smem_t = SharedStorage< + BLOCK_Q, NUM_MAX_SPARSE_BLOCKS, kNumKVBlocksPerSplit, kNumKVSplitsPerEntry, kNumThreads>; + + // Template checks + DG_STATIC_ASSERT(BLOCK_Q == 2 and kNumThreads == 256 and kNumSMs <= kNumThreads, "Unsupported metadata shape"); + DG_STATIC_ASSERT(SPLIT_KV % SPARSE_BLOCK_KV == 0 and SPLIT_KV % 128 == 0, "Invalid sparse split shape"); + DG_STATIC_ASSERT(NUM_MAX_SPARSE_BLOCKS % 4 == 0, "Sparse index rows must be 16-byte aligned"); + DG_STATIC_ASSERT(NUM_MAX_SPARSE_BLOCKS <= 4096, "Too many sparse KV blocks"); + DG_STATIC_ASSERT(kNumKVBlocksPerSplit <= kInvalidSparseSlot, "Sparse slot offset overflow"); + DG_STATIC_ASSERT(kNumKVSplitsPerEntry > 0, "Invalid number of KV splits per entry"); + DG_STATIC_ASSERT(not kIsPaged or PAGE_KV % SPARSE_BLOCK_KV == 0, "Invalid page shape"); + DG_STATIC_ASSERT(not kIsPaged or not kUseUnalignedKs, "Paged sparse MQA does not use ks"); + + const uint32_t thread_idx = threadIdx.x; + cudaGridDependencySynchronize(); + + const auto workspace_state = reinterpret_cast(workspace); + const auto q_block_infos = reinterpret_cast(workspace + sizeof(WorkspaceState)); + const auto kv_splits = reinterpret_cast*>(metadata + sizeof(MetadataHeader)); + + // Shared memory + extern __shared__ __align__(16) uint8_t storage[]; + auto& smem = *reinterpret_cast(storage); + + // Build Q-block metadata + uint32_t q_token_idx = blockIdx.x * (kIsPaged ? 1 : BLOCK_Q); + while (true) { + if (thread_idx == 0) { + smem.q_block.num_q_tokens = 0; + while (q_token_idx < num_q_tokens) { + uint32_t num_q_block_tokens = cute::min(BLOCK_Q, num_q_tokens - q_token_idx); + if constexpr (kIsPaged) { + // Keep Q blocks within one request + const uint32_t request_idx = indices[q_token_idx]; + // Requests contain few Q tokens, so use a linear scan + uint32_t request_q_token_base = q_token_idx; + while (request_q_token_base > 0 and indices[request_q_token_base - 1] == request_idx) + --request_q_token_base; + if ((q_token_idx - request_q_token_base) % BLOCK_Q != 0) { + q_token_idx = gridDim.x + ptx::atomic_add(&workspace_state->next_q_offset, 1u); + continue; + } + num_q_block_tokens = q_token_idx + 1 < num_q_tokens and + indices[q_token_idx + 1] == request_idx ? BLOCK_Q : 1; + } + const auto get_num_kv_blocks = [&](const uint32_t q_idx) { + const uint32_t kv_begin = kIsPaged ? 0 : cute::min(cu_seq_len_k_start[q_idx], num_kv_tokens); + const uint32_t kv_end = kIsPaged ? context_lens[q_idx] : + cute::max(kv_begin, cute::min(cu_seq_len_k_end[q_idx], num_kv_tokens)); + if constexpr (not kIsPaged and not kUseUnalignedKs) + DG_DEVICE_ASSERT(kv_end == kv_begin or kv_begin % SPARSE_BLOCK_KV == 0); + return cute::min(NUM_MAX_SPARSE_BLOCKS, math::ceil_div(kv_end - kv_begin, SPARSE_BLOCK_KV)); + }; + smem.q_block.q_token_base = q_token_idx; + smem.q_block.num_q_tokens = num_q_block_tokens; + smem.q_block.num_kv_blocks[0] = get_num_kv_blocks(q_token_idx); + smem.q_block.num_kv_blocks[1] = num_q_block_tokens == BLOCK_Q ? get_num_kv_blocks(q_token_idx + 1) : 0; + break; + } + } + __syncthreads(); + + const uint32_t num_q_block_tokens = smem.q_block.num_q_tokens; + if (num_q_block_tokens == 0) + break; + const uint32_t q_token_base = smem.q_block.q_token_base; + const uint32_t num_kv_blocks_in_q0 = smem.q_block.num_kv_blocks[0]; + const uint32_t num_kv_blocks_in_q1 = smem.q_block.num_kv_blocks[1]; + + // Load sparse indices + #pragma unroll + for (uint32_t q_token_offset = 0; q_token_offset < BLOCK_Q; ++ q_token_offset) { + const auto src = sparse_kv_block_indices + static_cast(q_token_base + q_token_offset) * NUM_MAX_SPARSE_BLOCKS; + const auto dst = smem.logical_kv_block_indices[q_token_offset]; + const uint32_t num_kv_blocks = q_token_offset == 0 ? num_kv_blocks_in_q0 : num_kv_blocks_in_q1; + if constexpr (kUseUnalignedKs) { + if (num_kv_blocks != 0) { + const uint32_t kv_offset = cu_seq_len_k_start[q_token_base + q_token_offset] % SPARSE_BLOCK_KV; + for (uint32_t q_slot_idx = thread_idx * 4; q_slot_idx < num_kv_blocks; q_slot_idx += kNumThreads * 4) { + const auto blocks = ptx::ld_evict_first(reinterpret_cast(src + q_slot_idx)); + ptx::st_shared(dst + q_slot_idx, + blocks.x * SPARSE_BLOCK_KV + kv_offset, blocks.y * SPARSE_BLOCK_KV + kv_offset, + blocks.z * SPARSE_BLOCK_KV + kv_offset, blocks.w * SPARSE_BLOCK_KV + kv_offset); + } + } + } else { + for (uint32_t q_slot_idx = thread_idx * 4; q_slot_idx < num_kv_blocks; q_slot_idx += kNumThreads * 4) + ptx::cp_async_cg<256>(reinterpret_cast(src + q_slot_idx), reinterpret_cast(dst + q_slot_idx)); + } + } + if constexpr (not kUseUnalignedKs) { + cutlass::arch::cp_async_fence(); + cutlass::arch::cp_async_wait<0>(); + } + __syncthreads(); + + // Merge-path partition + const uint32_t num_input_kv_blocks = num_kv_blocks_in_q0 + num_kv_blocks_in_q1; + const uint32_t merge_begin = thread_idx * num_input_kv_blocks / kNumThreads; + const uint32_t merge_end = (thread_idx + 1) * num_input_kv_blocks / kNumThreads; + uint32_t lo = cute::max(merge_begin, num_kv_blocks_in_q1) - num_kv_blocks_in_q1; + uint32_t hi = cute::min(merge_begin, num_kv_blocks_in_q0); + while (lo < hi) { + const uint32_t q0_slot_idx = (lo + hi) / 2; + const uint32_t q1_slot_idx = merge_begin - q0_slot_idx; + if (q1_slot_idx > 0 and q0_slot_idx < num_kv_blocks_in_q0 and + ptx::ld_shared(smem.logical_kv_block_indices[1] + q1_slot_idx - 1) >= ptx::ld_shared(smem.logical_kv_block_indices[0] + q0_slot_idx)) + lo = q0_slot_idx + 1; + else + hi = q0_slot_idx; + } + uint32_t q0_slot_idx = lo; + uint32_t q1_slot_idx = merge_begin - lo; + uint32_t num_remaining_inputs = merge_end - merge_begin; + uint32_t num_merged_kv_blocks_in_thread = 0; + + // Drop a duplicate carried across merge partitions + if (num_remaining_inputs > 0 and q0_slot_idx > 0 and q1_slot_idx < num_kv_blocks_in_q1 and + ptx::ld_shared(smem.logical_kv_block_indices[0] + q0_slot_idx - 1) == ptx::ld_shared(smem.logical_kv_block_indices[1] + q1_slot_idx)) { + ++ q1_slot_idx; + -- num_remaining_inputs; + } + + // Merge and pack each Q's sparse slot and presence bit + constexpr uint32_t kPresentBit = 1u << (kNumSparseSlotBits - 1); + constexpr uint32_t kSlotIndexMask = kPresentBit - 1; + const auto pack_slot = [](const uint32_t q_slot_idx, const bool is_present) { + return q_slot_idx | (is_present ? kPresentBit : 0u); + }; + constexpr uint32_t kNumKVBlocksPerThread = + math::constexpr_ceil_div(smem_t::kNumMaxMergedKVBlocks, kNumThreads); + uint32_t packed_slots_in_thread[kNumKVBlocksPerThread]; + #pragma unroll + for (uint32_t merged_kv_block_offset_in_thread = 0; merged_kv_block_offset_in_thread < kNumKVBlocksPerThread; ++ merged_kv_block_offset_in_thread) { + if (num_remaining_inputs == 0) + continue; + const uint32_t q0_logical_kv_block_idx = q0_slot_idx < num_kv_blocks_in_q0 ? ptx::ld_shared(smem.logical_kv_block_indices[0] + q0_slot_idx) : ~0u; + const uint32_t q1_logical_kv_block_idx = q1_slot_idx < num_kv_blocks_in_q1 ? ptx::ld_shared(smem.logical_kv_block_indices[1] + q1_slot_idx) : ~0u; + const bool in_q0 = q0_logical_kv_block_idx <= q1_logical_kv_block_idx; + const bool in_q1 = q1_logical_kv_block_idx <= q0_logical_kv_block_idx; + // Carry a final duplicate into the next partition + const bool consume_q1 = in_q1 and num_remaining_inputs > in_q0; + packed_slots_in_thread[merged_kv_block_offset_in_thread] = + pack_slot(q0_slot_idx, in_q0) | (pack_slot(q1_slot_idx, in_q1) << kNumSparseSlotBits); + ++ num_merged_kv_blocks_in_thread; + q0_slot_idx += in_q0; + q1_slot_idx += consume_q1; + num_remaining_inputs -= in_q0 + consume_q1; + } + DG_DEVICE_ASSERT(num_remaining_inputs == 0); + + // Compact merged blocks and reserve their KV-split range + uint32_t num_merged_kv_blocks; + const uint32_t merged_kv_block_base = math::cta_exclusive_sum( + num_merged_kv_blocks_in_thread, smem.warp_sums, num_merged_kv_blocks); + const uint32_t num_kv_splits_in_q_block = math::ceil_div(num_merged_kv_blocks, kNumKVBlocksPerSplit); + if (thread_idx == 0) { + smem.q_block.kv_split_base = static_cast(num_kv_splits_in_q_block == 0 ? 0 : + ptx::atomic_add(&workspace_state->num_kv_splits, num_kv_splits_in_q_block)); + q_block_infos[q_token_base] = QBlockInfo(smem.q_block.kv_split_base, num_kv_splits_in_q_block); + if constexpr (kIsPaged) { + if (num_q_block_tokens == BLOCK_Q) + q_block_infos[q_token_base + 1] = QBlockInfo(0, 0); + } + } + #pragma unroll + for (uint32_t merged_kv_block_offset_in_thread = 0; merged_kv_block_offset_in_thread < kNumKVBlocksPerThread; ++ merged_kv_block_offset_in_thread) { + if (merged_kv_block_offset_in_thread >= num_merged_kv_blocks_in_thread) + continue; + const uint32_t merged_kv_block_idx = merged_kv_block_base + merged_kv_block_offset_in_thread; + if constexpr (kIsPaged) { + if (merged_kv_block_idx % kNumKVBlocksPerSplit == 0) { + smem.packed_slot_bases_by_kv_split[merged_kv_block_idx / kNumKVBlocksPerSplit] = packed_slots_in_thread[merged_kv_block_offset_in_thread] & + (kSlotIndexMask | (kSlotIndexMask << kNumSparseSlotBits)); + } + } else { + smem.packed_slots_by_merged_kv_block[merged_kv_block_idx] = packed_slots_in_thread[merged_kv_block_offset_in_thread]; + } + } + __syncthreads(); + + const uint32_t kv_split_base = smem.q_block.kv_split_base; + const auto write_merged_kv_block = [&](const uint32_t merged_kv_block_idx, const uint32_t packed_slots) { + const uint32_t kv_split_offset = merged_kv_block_idx / kNumKVBlocksPerSplit; + const uint32_t kv_split_idx = kv_split_base + kv_split_offset; + const uint32_t kv_block_idx_in_split = merged_kv_block_idx % kNumKVBlocksPerSplit; + // Paged stores one packed Q-slot base per split; non-paged retains each merged KV block's Q slots for its second pass + uint32_t packed_slot_bases; + if constexpr (kIsPaged) + packed_slot_bases = ptx::ld_shared(smem.packed_slot_bases_by_kv_split + kv_split_offset); + else + packed_slot_bases = ptx::ld_shared(smem.packed_slots_by_merged_kv_block + merged_kv_block_idx - kv_block_idx_in_split); + const uint32_t q0_slot_idx = packed_slots & kSlotIndexMask; + const uint32_t q1_slot_idx = (packed_slots >> kNumSparseSlotBits) & kSlotIndexMask; + const uint32_t q0_slot_base = packed_slot_bases & kSlotIndexMask; + const uint32_t q1_slot_base = (packed_slot_bases >> kNumSparseSlotBits) & kSlotIndexMask; + const bool in_q0 = (packed_slots & kPresentBit) != 0; + const bool in_q1 = ((packed_slots >> kNumSparseSlotBits) & kPresentBit) != 0; + const uint32_t logical_kv_block_idx = in_q0 ? ptx::ld_shared(smem.logical_kv_block_indices[0] + q0_slot_idx) + : ptx::ld_shared(smem.logical_kv_block_indices[1] + q1_slot_idx); + uint32_t physical_kv_block_idx; + if constexpr (kIsPaged) { + constexpr uint32_t kNumKVBlocksPerPage = PAGE_KV / SPARSE_BLOCK_KV; + const uint32_t logical_page_idx = logical_kv_block_idx / kNumKVBlocksPerPage; + DG_DEVICE_ASSERT(logical_page_idx < block_table_stride); + const uint64_t block_table_idx = static_cast(q_token_base) * block_table_stride + logical_page_idx; + physical_kv_block_idx = block_table[block_table_idx] * kNumKVBlocksPerPage + logical_kv_block_idx % kNumKVBlocksPerPage; + } else { + physical_kv_block_idx = logical_kv_block_idx * (kUseUnalignedKs ? 1 : SPARSE_BLOCK_KV); + } + DG_DEVICE_ASSERT(not in_q0 or q0_slot_idx - q0_slot_base < kInvalidSparseSlot); + DG_DEVICE_ASSERT(not in_q1 or q1_slot_idx - q1_slot_base < kInvalidSparseSlot); + const uint32_t q0_slot_offset = in_q0 ? q0_slot_idx - q0_slot_base : kInvalidSparseSlot; + const uint32_t q1_slot_offset = in_q1 ? q1_slot_idx - q1_slot_base : kInvalidSparseSlot; + if (kv_block_idx_in_split == 0) { + const uint32_t num_kv_blocks_in_split = cute::min(kNumKVBlocksPerSplit, num_merged_kv_blocks - merged_kv_block_idx); + bool is_contiguous = false; + if constexpr (not kIsPaged and not kUseUnalignedKs) { + // TMA copies full splits; partial splits stay on the generic path + const uint32_t last_packed_slots = ptx::ld_shared(smem.packed_slots_by_merged_kv_block + merged_kv_block_idx + num_kv_blocks_in_split - 1); + const bool last_in_q0 = (last_packed_slots & kPresentBit) != 0; + const uint32_t last_q_slot_idx = last_in_q0 ? last_packed_slots & kSlotIndexMask + : (last_packed_slots >> kNumSparseSlotBits) & kSlotIndexMask; + const auto last_logical_kv_block_indices = smem.logical_kv_block_indices[last_in_q0 ? 0 : 1]; + const uint32_t last_logical_kv_block_idx = ptx::ld_shared(last_logical_kv_block_indices + last_q_slot_idx); + is_contiguous = num_kv_blocks_in_split == kNumKVBlocksPerSplit and + last_logical_kv_block_idx == logical_kv_block_idx + num_kv_blocks_in_split - 1; + } + kv_splits[kv_split_idx].header = KVSplitHeader(q_token_base, num_kv_blocks_in_split, is_contiguous, q0_slot_base, + num_q_block_tokens == BLOCK_Q ? q1_slot_base : ~0u); + } + kv_splits[kv_split_idx].kv_block_infos[kv_block_idx_in_split] = KVBlockInfo(physical_kv_block_idx, q0_slot_offset, q1_slot_offset); + }; + + // Write KV split metadata + if constexpr (kIsPaged) { + #pragma unroll + for (uint32_t merged_kv_block_offset_in_thread = 0; merged_kv_block_offset_in_thread < kNumKVBlocksPerThread; ++ merged_kv_block_offset_in_thread) { + if (merged_kv_block_offset_in_thread >= num_merged_kv_blocks_in_thread) + continue; + const uint32_t merged_kv_block_idx = merged_kv_block_base + merged_kv_block_offset_in_thread; + write_merged_kv_block(merged_kv_block_idx, packed_slots_in_thread[merged_kv_block_offset_in_thread]); + } + } else { + for (uint32_t merged_kv_block_idx = thread_idx; merged_kv_block_idx < num_merged_kv_blocks; merged_kv_block_idx += kNumThreads) + write_merged_kv_block(merged_kv_block_idx, ptx::ld_shared(smem.packed_slots_by_merged_kv_block + merged_kv_block_idx)); + } + + // Pad the last split to keep the main-kernel copy loop branch-free + for (uint32_t padded_kv_block_slot_idx = num_merged_kv_blocks + thread_idx; padded_kv_block_slot_idx < num_kv_splits_in_q_block * kNumKVBlocksPerSplit; + padded_kv_block_slot_idx += kNumThreads) { + kv_splits[kv_split_base + padded_kv_block_slot_idx / kNumKVBlocksPerSplit].kv_block_infos[padded_kv_block_slot_idx % kNumKVBlocksPerSplit] = + KVBlockInfo(0, kInvalidSparseSlot, kInvalidSparseSlot); + } + if (thread_idx == 0) { + q_token_idx = kIsPaged ? gridDim.x + ptx::atomic_add(&workspace_state->next_q_offset, 1u) : + q_token_idx + gridDim.x * BLOCK_Q; + } + } + + // The last CTA acquires all producer writes and builds the schedule + if (thread_idx == 0) { + smem.is_last_cta = ptx::atomic_add_rel(&workspace_state->num_finished_ctas, 1u) + 1 == gridDim.x; + if (smem.is_last_cta != 0) + ptx::ld_acq(&workspace_state->num_finished_ctas); + } + __syncthreads(); + if (smem.is_last_cta == 0) + return; + + const uint32_t total_kv_splits = workspace_state->num_kv_splits; + const auto schedule_entries = reinterpret_cast(kv_splits + total_kv_splits); + uint32_t num_waves; + if constexpr (not kIsPaged) { + num_waves = build_contiguous_schedule( + schedule_entries, kv_splits, q_block_infos, num_q_tokens, total_kv_splits, smem); + } else { + num_waves = build_paged_schedule( + schedule_entries, q_block_infos, num_q_tokens, indices, smem); + } + if (thread_idx == 0) { + const auto header = reinterpret_cast(metadata); + header->num_kv_splits = total_kv_splits; + header->num_waves = num_waves; + header->use_unaligned_ks = kUseUnalignedKs; + workspace_state->num_kv_splits = 0; + workspace_state->next_q_offset = 0; + workspace_state->num_finished_ctas = 0; + } +} + +} // namespace deep_gemm::sched::sparse_mqa_logits diff --git a/deep_gemm/legacy/m_grouped_gemm.py b/deep_gemm/legacy/m_grouped_gemm.py index e685a9ab01..d27d38933e 100644 --- a/deep_gemm/legacy/m_grouped_gemm.py +++ b/deep_gemm/legacy/m_grouped_gemm.py @@ -62,7 +62,7 @@ def m_grouped_bf16_gemm_nt_contiguous_tl(a: torch.Tensor, b: torch.Tensor, d: to m_indices: torch.Tensor): r0, r1, r2 = b.shape - assert a.is_contiguous() and (b.is_contiguous or b.mT.is_contiguous()) + assert a.is_contiguous() and (b.is_contiguous() or b.mT.is_contiguous()) assert m_indices.is_contiguous() and d.is_contiguous() assert a.dtype == torch.bfloat16 and b.dtype == torch.bfloat16 assert m_indices.dtype == torch.int32 and d.dtype == torch.bfloat16 diff --git a/deep_gemm/mega/__init__.py b/deep_gemm/mega/__init__.py index 2314564e22..10cb51761c 100644 --- a/deep_gemm/mega/__init__.py +++ b/deep_gemm/mega/__init__.py @@ -22,7 +22,12 @@ def __init__(self, group: dist.ProcessGroup, hidden: int, intermediate_hidden: int, num_shared_experts: int = 0, mma_type: str = 'fp8xfp4', - activation: str = 'swiglu'): + activation: str = 'swiglu', + base: Optional['SymmBuffer'] = None): + # Align token count + num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe()) + + # Init assert activation == 'swiglu', f'Only `swiglu` activation is supported, got `{activation}`' self.group = group self.num_experts = num_experts @@ -30,8 +35,11 @@ def __init__(self, group: dist.ProcessGroup, self.num_topk = num_topk self.hidden = hidden self.intermediate_hidden = intermediate_hidden + self.num_shared_experts = num_shared_experts + self.mma_type = mma_type + self.activation = activation - # Allocate a symmetric buffer + # Allocate or reuse a symmetric buffer num_bytes, slice_input_buffers = _C.get_symm_buffer_size_for_mega_moe( group.size(), num_experts, num_max_tokens_per_rank, num_topk, @@ -39,16 +47,25 @@ def __init__(self, group: dist.ProcessGroup, mma_type, activation, num_shared_experts ) - allocator = torch if group.size() == 1 else symm_mem - self.buffer = allocator.empty(num_bytes, dtype=torch.int8, device='cuda') - self.handle = ( - types.SimpleNamespace(buffer_ptrs=[self.buffer.data_ptr()]) - if group.size() == 1 - else symm_mem.rendezvous(self.buffer, group=group) - ) - self.buffer.zero_() - self.group.barrier() - torch.cuda.synchronize() + if base is None: + allocator = torch if group.size() == 1 else symm_mem + self.buffer = allocator.empty(num_bytes, dtype=torch.int8, device='cuda') + self.handle = ( + types.SimpleNamespace(buffer_ptrs=[self.buffer.data_ptr()]) + if group.size() == 1 + else symm_mem.rendezvous(self.buffer, group=group) + ) + self.buffer.zero_() + self.group.barrier() + torch.cuda.synchronize() + else: + assert base.buffer is not None and base.handle is not None and base.group is group, \ + 'Cannot reuse an invalid symmetric buffer' + assert num_bytes <= base.buffer.nbytes, \ + (f'The reused Mega MoE config requires {num_bytes} bytes, ' + f'but the symmetric buffer only has {base.buffer.nbytes} bytes') + self.buffer = base.buffer + self.handle = base.handle # Create input buffer views (self.x, self.x_sf, @@ -66,6 +83,7 @@ def destroy(self): self.x_sf = None +# TODO: remove this function def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, num_experts: int, num_max_tokens_per_rank: int, num_topk: int, @@ -74,9 +92,6 @@ def get_symm_buffer_for_mega_moe(group: dist.ProcessGroup, use_fp8_dispatch: Union[bool, None] = None, mma_type: str = 'fp8xfp4', activation: str = 'swiglu') -> SymmBuffer: - # Align token count - num_max_tokens_per_rank = align(num_max_tokens_per_rank, _C.get_token_alignment_for_mega_moe()) - # Backward compat: derive `mma_type` from `use_fp8_dispatch` if provided if use_fp8_dispatch is not None: assert use_fp8_dispatch == (mma_type.split('x')[0] == 'fp8') @@ -136,7 +151,7 @@ def transform_weights_for_mega_moe( Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]]: assert activation == 'swiglu', f'Only `swiglu` activation is supported, got `{activation}`' if isinstance(l1_weights, tuple): - # FP8: interleave gate/up for weight and SF, then transpose L1 SF for UTCCP + # FP8/FP4: interleave gate/up for weight and SF, then transpose L1 SF for UTCCP l1_w = _interleave_weights(l1_weights[0]) l1_sf = _transpose_sf_for_utccp(_interleave_weights(l1_weights[1])) l1_transformed = (l1_w, l1_sf) diff --git a/deep_gemm/utils/layout.py b/deep_gemm/utils/layout.py index 6512c5ab7a..d2eceae495 100644 --- a/deep_gemm/utils/layout.py +++ b/deep_gemm/utils/layout.py @@ -1,16 +1,8 @@ -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 - -# Valid for all CUDA versions 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, set_mk_alignment_for_contiguous_layout, get_mk_alignment_for_contiguous_layout, get_theoretical_mk_alignment_for_contiguous_layout, diff --git a/deep_gemm/utils/math.py b/deep_gemm/utils/math.py index a3caf5b051..07ab72326d 100644 --- a/deep_gemm/utils/math.py +++ b/deep_gemm/utils/math.py @@ -86,8 +86,11 @@ def _quantize_to_fp4_e2m1(x: torch.Tensor) -> torch.Tensor: # {0, 0.5, 1, 1.5, 2, 3, 4, 6} # midpoints: 0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0 code = torch.zeros_like(x, dtype=torch.uint8) - for boundary in (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0): - code += (ax > boundary).to(torch.uint8) + for boundary, round_up in zip( + (0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0), + (False, True, False, True, False, True, False), + ): + code += (ax >= boundary if round_up else ax > boundary).to(torch.uint8) sign = (x < 0) & (code != 0) code = code | (sign.to(torch.uint8) << 3) return code.view(torch.int8) diff --git a/docs/scaling-factor-format.md b/docs/scaling-factor-format.md new file mode 100644 index 0000000000..74d75755e0 --- /dev/null +++ b/docs/scaling-factor-format.md @@ -0,0 +1,469 @@ +# DeepGEMM Scaling Factor Format + +This document specifies the scaling factor (SF) tensor contract for DeepGEMM's FP8/FP4 GEMM APIs: required shapes, dtypes, strides, the transform pipeline, and how to call the APIs correctly. Source of truth: + +- `csrc/apis/layout.hpp` — SF transform dispatch (`transform_sf_into_required_layout`, `transform_k_grouped_sf_into_required_layout`) +- `csrc/utils/layout.hpp` — SF layout validation (`check_sf_layout`) +- `csrc/apis/gemm.hpp` — GEMM argument checks +- `csrc/apis/attention.hpp`, `csrc/apis/mega_moe.hpp`, `deep_gemm/mega/__init__.py` — attention / mega MoE SF requirements (Section 6) +- `csrc/jit_kernels/impls/smxx_layout.hpp` — host-side transform launchers +- `deep_gemm/include/deep_gemm/impls/smxx_layout.cuh` — transform kernels +- `tests/test_layout.py`, `tests/test_fp8_fp4.py`, `tests/generators.py` — working usage examples + +## 1. Definitions + +### 1.1 Quantization Recipe + +`recipe = (gran_m, gran_n, gran_k)` describes the **SF storage** granularity (not the compute granularity): one SF value covers a `gran_m x gran_k` block of A (or `gran_n x gran_k` of B). + +Given A of shape `[·, M, K]` and B of shape `[·, N, K]` (the leading `·` batch/group dimension is optional), the untransformed SFs must have: + +``` +SFA shape: [·, ceil_div(M, gran_m), ceil_div(K, gran_k)] dtype: float32 +SFB shape: [·, ceil_div(N, gran_n), ceil_div(K, gran_k)] dtype: float32 +``` + +Rules: + +- Pass exactly one of `recipe` or the `recipe_a` + `recipe_b` pair. Use `recipe_a = (gran_m, gran_k_a)` / `recipe_b = (gran_n, gran_k_b)` (2-tuples) when A and B use different K granularities. +- Supported `gran_k`: `32` or `128` on SM100; `128` only on SM90. + +### 1.2 SF Value Constraint + +Every `float32` SF value must be an exact power of 2: bit pattern `[0][8-bit exponent][23 mantissa bits = 0]`. Sign bit and mantissa must be zero; this is asserted on device (`DG_TRAP_ONLY_DEVICE_ASSERT((value & 0x807fffffu) == 0)` in `smxx_layout.cuh`). + +When producing SFs from a quantization cast kernel, pass `round_sf=True` to guarantee this. + +### 1.3 Constants + +- `packed_sf_dtype` = `int32` (4 UE8M0 exponents packed per element, see Section 7.1) +- `ALIGN_MN` = `16 bytes / sizeof(int32)` = `4` (TMA requires `stride(-1)` to be a multiple of 16 bytes) + +## 2. Two Ways to Provide SFs + +| Path | SF dtype | Extra kernel launch | When to use | +|---|---|---|---| +| A. Untransformed | `float32` | Yes — DeepGEMM launches a transform kernel per GEMM call | Prototyping, correctness testing | +| B. Pre-transformed | `int32` (packed UE8M0) | No — layout is only validated | Production; weights (transform once, cache) and activations (produce directly from the cast kernel) | + +For path B, the recipe passed to the GEMM must have `gran_m = gran_n = 1` (the transform broadcasts SFs along MN), with `gran_k` unchanged. + +## 3. Pre-transformed SF Format Contract (Path B) + +A pre-transformed SF tensor for `mn` rows and `k` columns with granularity `(1, gran_k)` must satisfy (validated by `check_sf_layout` in `csrc/utils/layout.hpp`): + +``` +dtype: int32 (packed UE8M0) +shape: [·, mn, packed_sf_k] where packed_sf_k = ceil_div(k, gran_k * 4) +stride(-3): stride(-1) * size(-1) # outer/group dimension packed tightly +stride(-2): 1 # contiguous along MN ("MN-major") +stride(-1): align(mn, 4) # TMA 16-byte alignment +``` + +Note the tensor is **MN-major**: the MN dimension is the contiguous one, so the underlying memory is `packed_sf_k` slices of `align(mn, 4)` elements each (see Section 7.2 for a diagram). + +## 4. The Transform: `transform_sf_into_required_layout` + +```python +deep_gemm.transform_sf_into_required_layout( + sf, # torch.Tensor, float32 or int32 + mn, # int: M (if is_sfa) or N + k, # int + recipe, # (gran_m, gran_n, gran_k) or (gran_mn, gran_k) + num_groups=None, # int: set if sf has a leading group dimension + is_sfa=None, # bool: REQUIRED with a 3-tuple recipe; FORBIDDEN with a 2-tuple + disable_ue8m0_cast=False, + psum_layout=None, # torch.Tensor: only for SFA under the PSUM layout (skips gap rows) +) -> torch.Tensor +``` + +Recipe form rules (asserted in `csrc/apis/layout.hpp`): + +- 3-tuple `(gran_m, gran_n, gran_k)`: must also pass `is_sfa` (`True` selects `gran_m`, `False` selects `gran_n`). +- 2-tuple `(gran_mn, gran_k)`: must NOT pass `is_sfa`. + +Dispatch table (`csrc/apis/layout.hpp`): + +| Input dtype | `gran_mn` | `gran_k` | Arch | Action | +|---|---|---|---|---| +| `float32` | 1 | 128 | SM90 (or `disable_ue8m0_cast`) | Transpose to MN-major, TMA-aligned `float32` (no packing) | +| `float32` | 128 | 128 | SM90 (or `disable_ue8m0_cast`) | Validate only (no transform) | +| `float32` | any | 32 or 128 | SM100 | Broadcast along MN to `gran_mn=1`, then pack to UE8M0 `int32`, MN-major, TMA-aligned | +| `int32` | 1 | 32 or 128 | SM100 | Validate only (already pre-transformed; must satisfy Section 3) | + +For the SM100 `float32` row, the returned tensor is exactly the pre-transformed format defined in Section 3: shape `[·, mn, ceil_div(k, gran_k * 4)]`, dtype `int32`, strides `(align(mn, 4) * ceil_div(k, gran_k * 4), 1, align(mn, 4))`. It can be cached and passed back to later calls, which then hit the `int32` validate-only row. + +### Example: transform weight SFs once and cache + +```python +import deep_gemm + +# Weights for fp8_einsum 'bhr,hdr->bhd': B operand = [h, d, r], so N=d, K=r, batch=h. +# Untransformed SFB: [h, ceil_div(d, 128), ceil_div(r, 128)], float32 +sfw = deep_gemm.transform_sf_into_required_layout( + sf=scale_factor, + mn=d, k=r, + recipe=(1, 128, 128), # (gran_m, gran_n, gran_k) of the ORIGINAL quantization + is_sfa=False, # this is SFB, so gran_mn = gran_n = 128 + num_groups=h, +) +# sfw: [h, d, ceil_div(r, 128 * 4)], int32, MN-major, TMA-aligned. Cache and reuse. + +deep_gemm.fp8_einsum( + 'bhr,hdr->bhd', + (x_fp8, sfx), + (w_fp8, sfw), # pre-transformed: no transform kernel launched for SFB + out, + recipe=(1, 1, 128), # gran_n is now 1 (broadcast during the transform); gran_k unchanged +) +``` + +This works for SFB because `'bhr,hdr->bhd'` does not permute the B operand. To also pre-transform the activation SFA, see Section 6.2 — `fp8_einsum` permutes SFs internally, which changes how the transform must be applied. + +### Example: produce pre-transformed SFs directly from a cast kernel + +Cast kernels can emit the packed `int32` SF directly, so no transform is ever needed (e.g., `per_token_cast` from `tile_kernels`): + +```python +from tile_kernels.quant import per_token_cast + +a_fp8, sfa = per_token_cast( + x=activation, # [num_tokens, hidden], bf16 + fmt='e4m3', + num_per_channels=128, # = gran_k + round_sf=True, # SF values are exact powers of 2 (Section 1.2) + use_tma_aligned_col_major_sf=True, # MN-major + TMA-aligned strides (Section 3) + use_packed_ue8m0=True, # packed int32 output +) +# a_fp8: [num_tokens, hidden], float8_e4m3fn +# sfa: [num_tokens, ceil_div(hidden, 128 * 4)], int32, satisfies Section 3 + +deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( + a=(a_fp8, sfa), # int32 SF: validated, not transformed + b=(b_fp4, sfb), + d=output, + grouped_layout=grouped_layout, + recipe_a=(1, 128), # (gran_m, gran_k): gran_m must be 1 for pre-transformed SFs + recipe_b=(1, 32), + use_psum_layout=True, +) +``` + +### Example: simplest path (untransformed float32 SFs) + +```python +# A: [m, k] FP8, quantized at 1x128; SFA: [m, ceil_div(k, 128)], float32 +# B: [num_groups, n, k] FP4, quantized at 32x32; SFB: [num_groups, ceil_div(n, 32), ceil_div(k, 32)], float32 +# grouped_layout: [num_groups], int32, PSUM row boundaries: +# group i occupies A rows [align(layout[i-1], alignment), layout[i]) +deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous( + a=(a_fp8, sfa), + b=(b_fp4, sfb), + d=output, + grouped_layout=grouped_layout, + recipe_a=(1, 128), + recipe_b=(32, 32), # float32 SFs may use any supported granularity + use_psum_layout=True, +) +# DeepGEMM launches the transform kernel internally for both SFs on every call. +``` + +## 5. K-Grouped Contiguous Layout + +K-grouped GEMM concatenates groups along K. The SM100 FP8 TN API uses logical operands `A = [sum_k, m]` and `B = [sum_k, n]`; the SM100 FP4 NT API uses `A = [m, sum_k]` and `B = [n, sum_k]`, backed by packed-byte storage `[m, sum_k / 2]` and `[n, sum_k / 2]`. Both produce `D = [num_groups, m, n]` and share the `grouped_layout` semantics below. APIs: + +```python +deep_gemm.k_grouped_fp8_gemm_tn_contiguous( # SM100; SM90 uses k_grouped_fp8_gemm_nt_contiguous + a, # (tensor, sf) pair + b, # (tensor, sf) pair + d, + ks_cpu, # List[int] or None: per-group K sizes, on CPU + grouped_layout, # torch.Tensor: [num_groups], int32, on device (num_groups <= 128) + c, # although the signature defaults to None, k-grouped asserts c is + # not None: pass an FP32 accumulator with the same shape as d + # (d = c + sum over groups); passing None fails a DG_HOST_ASSERT. + # Prefer passing the SAME tensor as d (c is d): the kernel then + # accumulates in place. If c is a different tensor, DeepGEMM first + # runs d.copy_(c) before the GEMM (an extra full copy) + recipe=(1, 1, 128), # gran_m and gran_n MUST be 1; gran_k in {32, 128} on SM100 + compiled_dims="mn", + use_psum_layout=False, +) + +deep_gemm.k_grouped_fp4_gemm_nt_contiguous( # SM100 only; same arguments + a, b, d, ks_cpu, grouped_layout, c, + recipe=(1, 1, 32), compiled_dims="mn", use_psum_layout=False, +) +``` + +`k_alignment` below is the global MK alignment for contiguous layouts, set via `deep_gemm.set_mk_alignment_for_contiguous_layout(value)`. It must be a multiple of the kernel's `BLOCK_K`: 128 for FP8 on SM100 (SM90: exactly 128) and 256 for FP4. Every aligned K range must have zero A/B padding and valid corresponding SF padding. + +### 5.1 `grouped_layout` and `ks_cpu` semantics + +The meaning of `grouped_layout` depends on `use_psum_layout`: + +| `use_psum_layout` | `grouped_layout[i]` contains | `ks_cpu` | `k_i` constraints | +|---|---|---|---| +| `False` | group `i`'s K size directly | required | each `k_i % k_alignment == 0` and `k_i % gran_k == 0` | +| `True` | cumulative (prefix-sum) end offset; group `i` occupies K range `[align(layout[i-1], k_alignment), layout[i])` | optional | `k_i` needs no alignment | + +With `use_psum_layout=True`: + +- If `ks_cpu` is provided, it must contain the **aligned** sizes `align(k_i, k_alignment)` (each entry must be a multiple of `k_alignment`); the exact SF shape is then computed on the host. +- If `ks_cpu` is `None` or `[]`, group sizes are read from `grouped_layout` on the device, and the host allocates an upper bound of `(sf_k + 3 * num_groups) / 4` packed rows. + +### 5.2 K-grouped SF contract + +The SF input (per operand) is 2D. Two accepted dtypes: + +**`float32` (DeepGEMM packs it):** shape `[sum_sf_k, mn]` contiguous, where `sum_sf_k` = total SF rows over all groups (`k_i / gran_k` rows per group without PSUM; `align(k_i, k_alignment) / gran_k` with PSUM). Supported for `gran_k` 32 and 128. Requires `mn % 4 == 0`. + +**`int32` pre-packed (validated only, no kernel):** ONLY accepted on SM100 with `gran_k = 32`. For `gran_k = 128` you must pass `float32`. Contract: + +``` +shape: [packed_sf_k, mn] where packed_sf_k >= sum(ceil_div(k_i, gran_k * 4)) # larger is OK, e.g. + a buffer pre-allocated for the maximum K; unused trailing rows are ignored +stride: [mn, 1] (contiguous); mn % 4 == 0 +layout: each group starts at a new packed row; when a group's SF row count is not a + multiple of 4, the trailing UE8M0 slots of its last packed row are zero-filled +``` + +Note the k-grouped packed SF is plain contiguous `[packed_sf_k, mn]` — unlike Section 3, no padding is inserted along MN; instead `mn % 4 == 0` is a hard precondition. + +### 5.3 Example: k-grouped with float32 SFs (non-PSUM) + +```python +import torch, deep_gemm + +gran_k, k_alignment = 32, 128 +deep_gemm.set_mk_alignment_for_contiguous_layout(k_alignment) + +ks = [2048, 4096, 1024] # each a multiple of k_alignment (and gran_k) +sum_k = sum(ks) +grouped_layout = torch.tensor(ks, device='cuda', dtype=torch.int32) # per-group sizes (non-PSUM) + +# a_fp8: [sum_k, m] float8_e4m3fn; sfa: [sum_k // gran_k, m] float32, contiguous +# b_fp8: [sum_k, n] float8_e4m3fn; sfb: [sum_k // gran_k, n] float32, contiguous +# m and n must be multiples of 4 (SF packing precondition, Section 5.2) +d = torch.zeros((num_groups, m, n), device='cuda', dtype=torch.float) +deep_gemm.k_grouped_fp8_gemm_tn_contiguous( + a=(a_fp8, sfa), b=(b_fp8, sfb), d=d, + ks_cpu=ks, grouped_layout=grouped_layout, + c=d, # same tensor as d: in-place accumulation, no extra copy + recipe=(1, 1, gran_k), +) +``` + +### 5.4 Example: k-grouped with pre-packed int32 SFs + +```python +# Pack once with the dedicated helper (or produce packed SFs from your cast kernel): +sfa_packed = deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + sfa, # [sum_sf_k, mn] float32, contiguous + grouped_layout, # semantics per Section 5.1 + ks_cpu=ks, # or None/[] with use_psum_layout=True + gran_k=32, # int32 pre-packing path requires gran_k == 32 + k_alignment=k_alignment, + use_psum_layout=False, +) +# sfa_packed: [sum(ceil_div(k_i, 32*4)), mn], int32 + +deep_gemm.k_grouped_fp8_gemm_tn_contiguous( + a=(a_fp8, sfa_packed), b=(b_fp8, sfb_packed), d=d, + ks_cpu=ks, grouped_layout=grouped_layout, c=d, + recipe=(1, 1, 32), # gran_k stays 32; layout is validated, no transform kernel runs +) +``` + +### 5.5 Example: PSUM layout (dynamic group sizes) + +```python +# real_ks may be unaligned; groups are stored padded to k_alignment +def build_psum_layout(real_ks, k_alignment): + psum, prev_end = [], 0 + for k in real_ks: + end = (prev_end + k_alignment - 1) // k_alignment * k_alignment + k # align(prev_end) + k + psum.append(end) + prev_end = end + return psum + +real_ks = [1000, 4096, 900] +grouped_layout = torch.tensor(build_psum_layout(real_ks, k_alignment), device='cuda', dtype=torch.int32) +aligned_ks = [(k + k_alignment - 1) // k_alignment * k_alignment for k in real_ks] + +deep_gemm.k_grouped_fp8_gemm_tn_contiguous( + a=(a_fp8, sfa), b=(b_fp8, sfb), d=d, + ks_cpu=aligned_ks, # pass ALIGNED sizes; or None if only known on device + grouped_layout=grouped_layout, + c=d, recipe=(1, 1, gran_k), + use_psum_layout=True, +) +``` + +## 6. SF Requirements Beyond Plain GEMM (Attention, Mega MoE) + +Several non-GEMM APIs take SF tensors with **hardcoded** requirements that bypass the recipe/transform pipeline of Section 4. Passing the wrong SF dtype fails a `DG_HOST_ASSERT` immediately. + +### 6.1 MQA Logits: `fp8_fp4_mqa_logits` / `fp8_fp4_paged_mqa_logits` + +Sources: `csrc/apis/attention.hpp` (host checks), `tests/test_attention.py` (SF construction). + +```python +deep_gemm.fp8_fp4_mqa_logits( + q, # (q_fp, q_sf or None): q_fp [seq_len, num_heads, head_dim] + kv, # (kv_fp, kv_sf): kv_fp [seq_len_kv, head_dim] + weights, # [seq_len, num_heads]; float32, or bf16 (SM100, forces bf16 logits) + cu_seq_len_k_start, cu_seq_len_k_end, + clean_logits=True, max_seqlen_k=0, logits_dtype=torch.float32, +) +``` + +**The dtype of `q_sf` and `kv_sf` is COUPLED.** Passing `q_sf` selects "MX mode", which flips the required `kv_sf` dtype (attention.hpp:123: `kv_sf.scalar_type() == (is_mx_sf ? kInt32 : kFloat)`). There is no mixed mode: + +| Mode | `q_sf` | `kv_sf` | Q/KV data dtype | Arch | +|---|---|---|---|---| +| MX (`q_sf` provided) | `int32` packed UE8M0, contiguous | `int32` packed UE8M0, contiguous | FP8 (MXFP8) or packed FP4 (MXFP4) | SM100 only | +| non-MX (`q_sf=None`) | — | `float32` (one plain scale per token), contiguous | FP8 only | SM90 / SM100 | + +Additional rules: + +- **FP4 Q/KV data requires MX mode** — `q_sf` must be provided (attention.hpp:92). +- SF shapes: `q_sf` is `[seq_len, num_heads]` (non-paged) or `[batch_size, next_n, num_heads]` (paged); `kv_sf` is 1-D `[seq_len_kv]`. Both contiguous. +- MX SF granularity is per-32-element blocks along `head_dim` — with `head_dim <= 128` all (up to 4) UE8M0 exponents of one token/head fit in exactly **one `int32`**, hence the shapes above have no trailing K dimension. +- The legacy aliases `fp8_mqa_logits` / `fp8_paged_mqa_logits` hardwire `q_sf=None`, i.e., always the non-MX `float32` mode. +- Paged variant: the KV SF is **fused into the byte cache**, not a separate tensor. `kv_cache` is `uint8` of shape `[num_kv_blocks, block_kv, 1, kv_head_dim + 4]` — per token, the value bytes (`head_dim` for FP8, `head_dim/2` for FP4) are followed by 4 SF bytes interpreted as `int32` (MX) or `float32` (non-MX) (attention.hpp:266-285). + +```python +from deep_gemm.utils import per_token_cast_to_fp8, per_custom_dims_cast_to_fp8 + +# MX mode (SM100): per-token 1x32 quantization, packed UE8M0 for BOTH SFs +q_fp8_2d, q_sf = per_token_cast_to_fp8(q.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) +kv_fp8, kv_sf = per_token_cast_to_fp8(kv, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) +logits = deep_gemm.fp8_fp4_mqa_logits( + q=(q_fp8_2d.view(seq_len, num_heads, head_dim), q_sf.view(seq_len, num_heads)), # int32 + kv=(kv_fp8, kv_sf.view(seq_len_kv)), # int32 + weights=weights, cu_seq_len_k_start=ks, cu_seq_len_k_end=ke) + +# non-MX mode: q_sf=None forces kv_sf to be plain float32 per-token +kv_fp8, kv_sf = per_custom_dims_cast_to_fp8(kv, (0,), False) # kv_sf: [seq_len_kv], float32 +logits = deep_gemm.fp8_fp4_mqa_logits( + q=(q.to(torch.float8_e4m3fn), None), + kv=(kv_fp8, kv_sf), + weights=weights, cu_seq_len_k_start=ks, cu_seq_len_k_end=ke) +``` + +### 6.2 `fp8_einsum`: SFs Are Permuted Internally + +`fp8_einsum` hardcodes its expressions and **permutes each operand AND its SF** into `(batch, m, n, k)` order before calling the internal batched GEMM (`csrc/apis/einsum.hpp:209-232`): + +| Expression | `(batch, m, n, k)` | SFA permute | SFB permute | +|---|---|---|---| +| `'bhr,hdr->bhd'` | `(h, b, d, r)` | `(1, 0, 2)` | none | +| `'bhd,hdr->bhr'` (SM100) | `(h, b, r, d)` | `(1, 0, 2)` | `(0, 2, 1)` | +| `'bhd,bhr->hdr'` (SM100) | `(h, d, r, b)` | `(1, 2, 0)` | `(1, 2, 0)` | + +Consequences for pre-transformed (`int32`) SFs — the Section 3 stride contract is checked on the **post-permute** tensor: + +- If the SF is not permuted (e.g., SFB of `'bhr,hdr->bhd'`, as in the Section 4 weight example), transform it directly with the operand's own `(num_groups, mn, k)`. +- If the SF is permuted, you must transform it **in post-permute coordinates** and permute it back before passing it to `fp8_einsum`. Do NOT `view`/`reshape` a 2-D transform output into 3-D — the transform output is MN-major (non-contiguous), so reshaping destroys the required strides. + +```python +# WRONG: 2-D transform + reshape breaks the MN-major strides +sfx_2d = deep_gemm.transform_sf_into_required_layout(sfx_f32.view(b * h, -1), mn=b * h, k=r, recipe=(1, 128)) +sfx_int32 = sfx_2d.reshape(b, h, -1) # silently copies to contiguous strides; + # stride(-2) != 1 after the internal permute -> DG_HOST_ASSERT fails + +# CORRECT for 'bhr,hdr->bhd' SFA: fp8_einsum permutes SFA with (1, 0, 2), so the internal +# batched GEMM sees [h, b, packed_sf_k]. Transform with num_groups=h, mn=b, then permute back: +sfx_grouped = deep_gemm.transform_sf_into_required_layout( + sfx_f32.permute(1, 0, 2).contiguous(), # [h, b, ceil_div(r, 128)], float32 + mn=b, k=r, recipe=(1, 128), num_groups=h, +) # [h, b, packed_sf_k], int32, Section 3 layout +sfx_int32 = sfx_grouped.permute(1, 0, 2) # [b, h, packed_sf_k] view; einsum permutes it back +deep_gemm.fp8_einsum('bhr,hdr->bhd', (x_fp8, sfx_int32), (w_fp8, sfw), out, recipe=(1, 1, 128)) +``` + +Untransformed `float32` SFs need no special care — the internal transform handles the permuted layout. + +### 6.3 Mega MoE: `fp8_fp4_mega_moe` + +Sources: `csrc/apis/mega_moe.hpp` (host checks), `deep_gemm/mega/__init__.py` (Python wrapper + weight transform), `tests/test_mega_moe.py` (SF construction). + +- **Recipe is pinned to `(1, 1, 32)`** (mega.hpp:179); `kGranK = 32` is also hardcoded in the JIT impl (sm100_fp8_fp4_mega_moe.hpp:165). +- **Weight SFs must already be packed `int32`** — unlike the GEMM APIs, there is NO float32 fallback and no auto-transform: `check_sf_layout(..., type_check=torch::kInt)` rejects `float32` outright (mega.hpp:206-209, 229-232). The required layout is exactly the Section 3 contract with `gran_k = 32`: + - Routed L1 SF: `[num_experts_per_rank, 2*intermediate_hidden, ceil_div(hidden, 128)]`, `int32`, MN-major, TMA-aligned (`stride(-2)==1`, `stride(-1)==align(mn,4)`); routed L2 SF: `[num_experts_per_rank, hidden, ceil_div(intermediate_hidden, 128)]`. (The `/128` is `gran_k * 4 = 32 * 4`.) + - Shared expert SFs (optional; weights are FP8 instead of FP4): same contract, 2-D without the group dimension. +- **An extra mega-MoE-only layout step is mandatory**: pass weights + SFs through `deep_gemm.transform_weights_for_mega_moe` before the call. It (a) interleaves gate/up rows of L1 (weights AND SF, 8-row granularity) and (b) applies a UTCCP intra-128-row transpose to both L1 and L2 SFs (`reshape(-1, 4, 32, packed_sf_k).transpose(2, 3)`; requires `mn % 128 == 0`) — `deep_gemm/mega/__init__.py:97-149`. The SFs in Section 3 layout alone are NOT directly consumable by the kernel. +- **Activation SFs are internal**: they live in slices of the symmetric buffer (`x_sf` is K-major `[num_max_tokens_per_rank, hidden // 128]` `int32`; the intermediate `l1/l2_acts_sf` are MN-major), produced by the kernel pipeline — not user-supplied arguments (mega.hpp:96-153). + +```python +from deep_gemm.utils import per_token_cast_to_fp4 + +# Routed expert weights: [g, n, k] bf16 -> FP4 (1x32) + float32 SF [g, n, k/32] +w = torch.empty((g, n, k // 2), device='cuda', dtype=torch.int8) +w_sf = torch.empty((g, n, k // 32), device='cuda', dtype=torch.float) +for i in range(g): + w[i], w_sf[i] = per_token_cast_to_fp4(w_bf16[i], use_ue8m0=True, gran_k=32) + +# Step 1: pack to the Section 3 int32 layout (mandatory; float32 SFs are rejected) +w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups=g) + +# Step 2: mega-MoE weight/SF shuffle (gate-up interleave + UTCCP SF transpose; mandatory) +(l1_w, l1_sf), (l2_w, l2_sf) = deep_gemm.transform_weights_for_mega_moe((l1_w, l1_sf), (l2_w, l2_sf)) +``` + +## 7. Internals + +Reference for kernel developers and for debugging layout mismatches. + +### 7.1 UE8M0 Packing + +The 8-bit exponents of 4 consecutive K positions are packed into one `int32`, little-endian by K index (`smxx_layout.cuh`): + +```cpp +uint32_t packed = 0; +packed |= (values[0] >> 23u); // exp of sf[4k+0] -> bits [7:0] +packed |= (values[1] >> 15u); // exp of sf[4k+1] -> bits [15:8] +packed |= (values[2] >> 7u); // exp of sf[4k+2] -> bits [23:16] +packed |= (values[3] << 1u); // exp of sf[4k+3] -> bits [31:24] +``` + +### 7.2 Memory Layout of the Transformed Tensor (non-k-grouped) + +``` +Shape: [mn, packed_sf_k] where packed_sf_k = ceil_div(k, gran_k * 4) +Stride: [1, align(mn, 4)] + +Diagram (mn=6, packed_sf_k=3, align(6,4)=8), int32 elements: + +Offset: 0 1 2 3 4 5 6 7 | 8 9 10 11 12 13 14 15 | 16 ... +Content: m0 m1 m2 m3 m4 m5 __ __ | m0 m1 m2 m3 m4 m5 __ __ | m0 ... + <--- K-slice 0 -----> <--- K-slice 1 -----> +``` + +Each K-slice occupies `align(mn, 4)` elements; the `__` padding exists only for 16-byte TMA alignment and is never read as data. + +### 7.3 PSUM Gap Row Handling (M-grouped) + +Under the M-grouped PSUM layout, gap rows exist between groups: + +``` +grouped_layout = [100, 250, 370], alignment = 128 +Group 0: rows [0, 100) valid +Gap: rows [100, 128) padding +Group 1: rows [128, 250) valid +Gap: rows [250, 256) padding +Group 2: rows [256, 370) valid +``` + +When `psum_layout` is passed to the SF transform, gap rows are not read from the input; the kernel writes `0` for them (a safe finite scale code — UE8M0 `0xff` is NaN). The GEMM kernel never consumes those values. + +### 7.4 K-Grouped Packing Algorithm + +Each group is packed independently, then concatenated along K (`pack_fp32_into_ue8m0` in `smxx_layout.cuh`): + +1. Determine group `i`'s input SF row range: + - Non-PSUM: `grouped_layout[i]` is the group K size `k_i` (a multiple of `k_alignment` and `gran_k`), covering `k_i / gran_k` rows. + - PSUM: `k_i = grouped_layout[i] - align(grouped_layout[i-1], k_alignment)`; the group covers its aligned region, `align(k_i, k_alignment) / gran_k` rows (PSUM data is stored padded to `k_alignment`). +2. Emit `ceil_div(num_group_sf_rows, 4)` packed `int32` rows for the group; if `num_group_sf_rows % 4 != 0`, zero-fill the trailing UE8M0 slots of the last packed row. +3. Concatenate all groups: output shape `[sum(ceil_div(num_group_sf_rows_i, 4)), mn]`, contiguous. diff --git a/setup.py b/setup.py index c4d74ae929..7b6deb1c51 100644 --- a/setup.py +++ b/setup.py @@ -22,13 +22,10 @@ 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', +cxx_flags = ['-std=c++20', '-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__)) @@ -37,10 +34,10 @@ f'{CUDA_HOME}/include', f'{CUDA_HOME}/include/cccl', 'deep_gemm/include', + 'third-party/deep_jit/include', 'third-party/cutlass/include', - 'third-party/fmt/include', ] -build_libraries = ['cudart', 'nvrtc'] +build_libraries = ['cudart'] build_library_dirs = [f'{CUDA_HOME}/lib64'] third_party_include_dirs = [ 'third-party/cutlass/include/cute', @@ -122,9 +119,20 @@ def run(self): # Third, generate and copy .pyi file to build root directory self.generate_pyi_file() + # Fourth, copy csrc and docs into the wheel for agent-side error lookup + self.prepare_agent_files() + # Finally, run the regular build build_py.run(self) + def prepare_agent_files(self): + # Copy csrc and docs into the wheel for agent-side error lookup + package_dir = os.path.join(self.build_lib, 'deep_gemm') + for name in ('csrc', 'docs'): + dst = os.path.join(package_dir, name) + shutil.rmtree(dst, ignore_errors=True) + shutil.copytree(os.path.join(current_dir, name), dst) + 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') diff --git a/tests/generators.py b/tests/generators.py index d7532ccf7a..7b336e1c94 100644 --- a/tests/generators.py +++ b/tests/generators.py @@ -2,6 +2,7 @@ import itertools import random import torch +from math import prod from typing import Generator, List, Optional, Tuple from deep_gemm.testing import get_arch_major @@ -53,6 +54,9 @@ def print(self): 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 is_fp4_fp4(self) -> bool: + return self.is_fp4_a and self.is_fp4_b + def get_recipes(self, is_wgrad: bool = False) -> Tuple[Tuple, Tuple, Tuple]: recipe, recipe_a, recipe_b = None, None, None if self.is_legacy(): @@ -73,9 +77,12 @@ def max_diff(self) -> float: def get_list_from_dtype(dtype: torch.dtype) -> List: if dtype == torch.bfloat16: return [None] + if dtype == torch.float4_e2m1fn_x2: + return [QuantConfig((32, 32, True, True))] quant_config_list = [QuantConfig()] if get_arch_major() == 10: quant_config_list.append(QuantConfig((128, 32, False, True))) + quant_config_list.append(QuantConfig((32, 32, True, True))) return quant_config_list @@ -112,8 +119,9 @@ def get_psum_layout_usage() -> tuple: return True, False -def enumerate_normal(dtype: torch.dtype) -> Generator: - assert dtype in (torch.float8_e4m3fn, torch.bfloat16) +def enumerate_normal(dtype: torch.dtype, collect_cublas_scores: bool = False) -> Generator: + assert dtype in (torch.float8_e4m3fn, torch.float4_e2m1fn_x2, torch.bfloat16) + assert not collect_cublas_scores or dtype != torch.bfloat16 quant_config_list = QuantConfig.get_list_from_dtype(dtype) fp32_output_nk = [(256, 7168), (129280, 7168)] @@ -127,31 +135,60 @@ def enumerate_normal(dtype: torch.dtype) -> Generator: for kernel_type in get_kernel_types(dtype): for quant_config in quant_config_list: + scores = [] if len(quant_config_list) > 1: quant_config.print() reset_seed() + def emit(*args): + return (*args, scores) if collect_cublas_scores else args + # Forward for m in m_fwd_list: - for i in range(len(nk_list)): - n, k = nk_list[i] + for i, (n, k) in enumerate(nk_list): 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 + yield emit(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 + yield emit(kernel_type, quant_config, m, n, k, MajorTypeAB.KMajor, MajorTypeAB.KMajor, True, 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 + if quant_config is None or not quant_config.is_fp4_fp4(): + 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 emit(kernel_type, quant_config, m, k, n, MajorTypeAB.KMajor, override_major, False, torch.bfloat16) # Dgrad + yield emit(override_kernel_type, quant_config, n, m, k, override_major, override_major, True, torch.float) # Wgrad + yield emit(override_kernel_type, quant_config, n, m, k, override_major, override_major, False, torch.bfloat16) # Wgrad + if dtype == torch.bfloat16 or get_arch_major() == 10: + yield emit(override_kernel_type, quant_config, n, m, k, override_major, override_major, False, torch.float) # Wgrad + + if collect_cublas_scores and scores: + quant_type = f'FP{4 if quant_config.is_fp4_a else 8}xFP{4 if quant_config.is_fp4_b else 8}' + print(f'Average {quant_type} GEMM speedup over cuBLASLt: ' + f'{prod(scores) ** (1.0 / len(scores)):.3f}x\n') + + +def enumerate_batched_syrk_symm() -> Generator: + # Real Muon parameter shapes from 5120/7168 models. Batch merging is capped at 16 in hai-llm. + b_m_k_list = [ + (16, 4608, 5120), + (16, 1536, 5120), + (16, 576, 5120), + (8, 24576, 1536), + (8, 32768, 512), + (4, 5120, 16384), + (8, 4096, 7168), + (8, 7168, 2048), + ] + for dtype in (torch.bfloat16, torch.float): + reset_seed() + for b, m, k in b_m_k_list: + yield b, m, k, dtype def enumerate_m_grouped_contiguous(dtype: torch.dtype) -> Generator: @@ -168,6 +205,8 @@ def enumerate_m_grouped_contiguous(dtype: torch.dtype) -> Generator: 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): + if quant_config is not None and quant_config.is_fp4_fp4() and major_b.is_mn_major(): + continue yield kernel_type, quant_config, num_groups, expected_m_per_group, n, k, major_a, major_b, use_psum_layout, ensure_zero_padding @@ -189,14 +228,25 @@ def enumerate_m_grouped_masked(dtype: torch.dtype) -> Generator: def enumerate_k_grouped_contiguous(dtype: torch.dtype): if dtype == torch.bfloat16: - k_alignment_options = [(128, 128)] if get_arch_major() == 9 else [(32, 32), (128, 128), (192, 192)] + sf_layout_list = [(128, 128)] if get_arch_major() == 9 else [(128, 128), (256, 256), (384, 384)] + elif dtype == torch.float4_e2m1fn_x2: + sf_layout_list = [(32, 256), (32, 512), (32, 768)] else: - k_alignment_options = [(128, 128)] if get_arch_major() == 9 else [(32, 32), (32, 128), (32, 160), (32, 224), (128, 128), (128, 160), (128, 224)] + sf_layout_list = ([(128, 128)] if get_arch_major() == 9 else + [(32, 128), (128, 128), (32, 256), (128, 256), (32, 384), (128, 384)]) # 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) + if get_arch_major() == 9 and dtype == torch.float8_e4m3fn: + major_a, major_b = MajorTypeAB.KMajor, MajorTypeAB.KMajor + elif dtype == torch.float4_e2m1fn_x2: + major_a, major_b = MajorTypeAB.KMajor, MajorTypeAB.KMajor + else: + major_a, major_b = MajorTypeAB.MNMajor, MajorTypeAB.MNMajor psum_list = (False, True) if get_arch_major() == 10 else (False, ) - # Must with FP32 accumulation and 1D1D kernels + if get_arch_major() == 9: + cd_options = [(True, torch.float)] + else: + cd_options = [(True, torch.float), (False, torch.float), (False, torch.bfloat16)] + # NOTES: the first shape has many small groups, for stressing the SM90 in-place tensor map update for num_groups, m, n, expected_k_per_group in (( 8, 768, 2048, 128), ( 4, 4096, 7168, 8192), ( 4, 7168, 2048, 8192), # EP64 @@ -204,28 +254,23 @@ def enumerate_k_grouped_contiguous(dtype: torch.dtype): (16, 4096, 7168, 2048), (16, 7168, 2048, 2048)): # EP16 real_ks_cpu = [max(1, int(expected_k_per_group * random.uniform(0.7, 1.3))) for _ in range(num_groups)] for use_psum_layout in psum_list: - for gran_k, k_alignment in k_alignment_options: + for gran_k, k_alignment in sf_layout_list: set_mk_alignment_for_contiguous_layout(k_alignment) aligned_ks_cpu = [align(k, k_alignment) for k in real_ks_cpu] - yield num_groups, m, n, major_a, major_b, real_ks_cpu, aligned_ks_cpu, expected_k_per_group, gran_k, k_alignment, use_psum_layout + for accumulate, out_dtype in cd_options: + yield (num_groups, m, n, major_a, major_b, real_ks_cpu, aligned_ks_cpu, + expected_k_per_group, gran_k, k_alignment, use_psum_layout, accumulate, out_dtype) -def enumerate_k_grouped_contiguous_test_variants(real_ks_cpu: List[int], k_alignment: int, - use_psum_layout: bool, include_k_tail: bool = False): - test_variants = [(False, False), (True, False)] - if include_k_tail: - test_variants.append((False, True)) +def enumerate_k_grouped_contiguous_test_variants(real_ks_cpu: List[int]): + yield list(real_ks_cpu) - for test_empty_groups, test_k_tail in test_variants: - test_real_ks_cpu = list(real_ks_cpu) - if test_empty_groups and len(real_ks_cpu) > 1: - test_real_ks_cpu[random.randint(0, len(real_ks_cpu) - 1)] = 0 - if test_k_tail and len(test_real_ks_cpu) > 1: - test_real_ks_cpu[0] = k_alignment - elif use_psum_layout and not test_empty_groups and len(test_real_ks_cpu) > 1 and test_real_ks_cpu[0] > 1: - test_real_ks_cpu[0] -= random.randint(1, min(k_alignment - 1, test_real_ks_cpu[0] - 1)) - test_aligned_ks_cpu = [align(k, k_alignment) for k in test_real_ks_cpu] - yield test_real_ks_cpu, test_aligned_ks_cpu, test_empty_groups, test_k_tail + empty_ks_cpu = [0 if random.random() < 0.25 else k for k in real_ks_cpu] + empty_ks_cpu[0], empty_ks_cpu[-1] = 0, real_ks_cpu[-1] + yield empty_ks_cpu + + single_group_ks_cpu = [k if i == len(real_ks_cpu) // 2 else 0 for i, k in enumerate(real_ks_cpu)] + yield single_group_ks_cpu def enumerate_sf_layout(): @@ -239,26 +284,28 @@ def enumerate_sf_layout(): set_mk_alignment_for_contiguous_layout(gran_k) yield mn, k, with_transpose, use_ue8m0, num_groups, gran_k + if get_arch_major() == 10: + for sf_k, use_ue8m0 in ((908, False), (1211, True)): + set_mk_alignment_for_contiguous_layout(32) + yield 4096, sf_k * 32, True, use_ue8m0, 1, 32 + 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_cpu = [align(int(random.uniform(0.7, 1.3) * avg_k), gran_k) for _ in range(num_groups)] - yield mn, ks_cpu, num_groups, gran_k + for gran_k, k_alignment in ((32, 128), (128, 128), (32, 256), (128, 256), (32, 384), (128, 384)): + set_mk_alignment_for_contiguous_layout(k_alignment) + ks_cpu = [align(int(random.uniform(0.7, 1.3) * avg_k), k_alignment) + for _ in range(num_groups)] + yield mn, ks_cpu, num_groups, gran_k, k_alignment def enumerate_k_grouped_psum_sf_layout(): - for mn, ks_cpu, num_groups, gran_k in enumerate_k_grouped_sf_layout(): - k_alignment_list = (32, 128, 160, 224) - for k_alignment in k_alignment_list: - # Generate non-aligned K sizes to test ceil_div path in pack kernel. - real_ks_cpu = [k + (gran_k // 2 if i % 2 else 0) for i, k in enumerate(ks_cpu)] - psum_layout = build_psum_layout_from_ks(real_ks_cpu, k_alignment) - aligned_ks_cpu = [align(k, k_alignment) for k in real_ks_cpu] - yield mn, real_ks_cpu, aligned_ks_cpu, psum_layout, num_groups, gran_k, k_alignment + for mn, ks_cpu, num_groups, gran_k, k_alignment in enumerate_k_grouped_sf_layout(): + real_ks_cpu = [k - (gran_k // 2 if i % 2 else 0) for i, k in enumerate(ks_cpu)] + aligned_ks_cpu = [align(k, k_alignment) for k in real_ks_cpu] + psum_layout = build_psum_layout_from_ks(real_ks_cpu, k_alignment) + yield mn, real_ks_cpu, aligned_ks_cpu, psum_layout, num_groups, gran_k, k_alignment def enumerate_transpose(): @@ -305,13 +352,15 @@ def generate_normal(m: int, n: int, k: int, accumulate: bool, out_dtype: torch.dtype, kernel_type: KernelType, use_ue8m0: bool = False, use_bf16: bool = False, - quant_config: Optional[QuantConfig] = None): + quant_config: Optional[QuantConfig] = None, + alpha: Optional[float] = 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) + alpha_value = alpha if alpha is not None else 1.0 + ref_d = (alpha_value * (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 @@ -341,6 +390,7 @@ def generate_m_grouped_contiguous(num_groups: int, expected_m_per_group: int, n: 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) + valid_mask = torch.zeros(m, device='cuda', dtype=torch.bool) start = 0 for i, (actual_m, aligned_m) in enumerate(zip(actual_ms, aligned_ms)): @@ -351,6 +401,7 @@ def generate_m_grouped_contiguous(num_groups: int, expected_m_per_group: int, n: else: grouped_layout[start: actual_end] = i grouped_layout[actual_end: aligned_end] = -1 + valid_mask[start:actual_end] = True # Zero BF16 padding so quantized SFA padding is regular, never uninitialized a[actual_end: aligned_end] = 0 ref_d[start: aligned_end] = a[start: aligned_end] @ b[i].t() @@ -358,20 +409,21 @@ def generate_m_grouped_contiguous(num_groups: int, expected_m_per_group: int, n: 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 + return m, a, b, grouped_layout, d, ref_d, valid_mask 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 + return m, a, b, grouped_layout, d, ref_d, valid_mask def layout_masked_to_psum(x: torch.Tensor, psum_m: torch.Tensor): - num_groups, max_m, _ = x.size() + num_groups, _, k = x.size() # PSUM gaps are intentionally left uninitialized to verify the pack kernel skips them - x_psum = torch.empty_like(x).view(num_groups * max_m, -1) + x_psum = torch.empty((align(psum_m[-1].item(), get_mk_alignment_for_contiguous_layout()), k), + dtype=x.dtype, device=x.device) 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] @@ -379,6 +431,18 @@ def layout_masked_to_psum(x: torch.Tensor, psum_m: torch.Tensor): return x_psum +def masked_valid_mask_to_psum(masked_m: torch.Tensor, psum_m: torch.Tensor) -> torch.Tensor: + # Build the PSUM-layout valid mask directly: gaps must be `False` (never uninitialized), + # otherwise selecting `d[valid_mask]` would read skipped rows and break determinism checks + total_m = align(psum_m[-1].item(), get_mk_alignment_for_contiguous_layout()) + valid_mask = torch.zeros(total_m, device=masked_m.device, dtype=torch.bool) + last_psum_m = 0 + for i in range(masked_m.numel()): + valid_mask[last_psum_m: psum_m[i]] = True + last_psum_m = align(psum_m[i], get_mk_alignment_for_contiguous_layout()) + return valid_mask + + 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, @@ -394,9 +458,15 @@ def generate_m_grouped_masked(num_groups: int, max_m: int, expected_m_per_group: 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 + valid_mask = torch.arange(max_m, device='cuda')[None, :] < masked_m[:, None] if use_bf16: - return a, b, masked_m, psum_m, d, ref_d + if use_psum_layout: + a = layout_masked_to_psum(a, psum_m) + d = layout_masked_to_psum(d, psum_m) + ref_d = layout_masked_to_psum(ref_d, psum_m) + valid_mask = masked_valid_mask_to_psum(masked_m, psum_m) + return a, b, psum_m if use_psum_layout else masked_m, d, ref_d, valid_mask 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) @@ -406,81 +476,70 @@ def generate_m_grouped_masked(num_groups: int, max_m: int, expected_m_per_group: # Zero SFA padding rows (beyond `masked_m`) so the pack kernel reads regular zeros for j in range(num_groups): a[1][j, masked_m[j].item():] = 0 + else: + a = (layout_masked_to_psum(a[0], psum_m), layout_masked_to_psum(a[1], psum_m)) + d = layout_masked_to_psum(d, psum_m) + ref_d = layout_masked_to_psum(ref_d, psum_m) + valid_mask = masked_valid_mask_to_psum(masked_m, psum_m) - return a, b, masked_m, psum_m, d, ref_d + return a, b, psum_m if use_psum_layout else masked_m, d, ref_d, valid_mask -def k_grouped_per_channel_cast_to_fp8(x: torch.Tensor, ks_cpu: List[int], use_ue8m0: bool, gran_k: int, - group_ends: Optional[List[int]] = None) -> Tuple[torch.Tensor, torch.Tensor]: - # Cast each group independently so that the SF rows stay compact (`ceil_div(k, gran_k)` rows per group); - # `group_ends` gives per-group end offsets (psum layouts), `None` means a contiguous prefix-sum layout +def k_grouped_cast_fp8_fp4_with_major(x: torch.Tensor, ks_cpu: List[int], major: MajorTypeAB, + use_ue8m0: bool, gran_k: int, is_fp4: bool, + group_ends: Optional[List[int]] = None, + sf_ks_cpu: Optional[List[int]] = None) -> Tuple[torch.Tensor, torch.Tensor]: assert x.dim() == 2 if group_ends is None: group_ends = list(itertools.accumulate(ks_cpu)) assert (group_ends[-1] if ks_cpu else 0) == x.size(0) + if sf_ks_cpu is None: + sf_ks_cpu = ks_cpu + assert len(ks_cpu) == len(group_ends) == len(sf_ks_cpu) - n = x.size(1) + mn = x.size(1) + if is_fp4: + assert major.is_k_major() and gran_k == 32 + data = torch.zeros((mn, x.size(0) // 2), dtype=torch.int8, device=x.device) + sf_groups = [] + for k, sf_k, end in zip(ks_cpu, sf_ks_cpu, group_ends): + if sf_k == 0: + continue + start = end - k + x_group = torch.zeros((sf_k, mn), dtype=x.dtype, device=x.device) + x_group[:k] = x[start:end] + x_group_fp4, x_group_sf = per_token_cast_to_fp4( + x_group.T.contiguous(), use_ue8m0=use_ue8m0, gran_k=gran_k) + data[:, start // 2:(start + sf_k) // 2] = x_group_fp4 + sf_groups.append(x_group_sf.T.contiguous()) + sf = torch.cat(sf_groups) if sf_groups else torch.empty((0, mn), dtype=torch.float, device=x.device) + return data, sf + + # Cast each group independently. SF rows are compact by group and padded to + # k_alignment, matching the K-grouped pack kernel's input contract. x_fp8 = torch.zeros(x.shape, dtype=torch.float8_e4m3fn, device=x.device) sf_groups = [] - for k, end in zip(ks_cpu, group_ends): - if k == 0: + for k, sf_k, end in zip(ks_cpu, sf_ks_cpu, group_ends): + if sf_k == 0: continue start = end - k - x_group = torch.zeros((align(k, gran_k), n), dtype=x.dtype, device=x.device) + x_group = torch.zeros((sf_k, mn), dtype=x.dtype, device=x.device) x_group[:k] = x[start:end] x_group_fp8, x_group_sf = per_channel_cast_to_fp8(x_group, use_ue8m0=use_ue8m0, gran_k=gran_k) x_fp8[start:end] = x_group_fp8[:k] sf_groups.append(x_group_sf) - sf = torch.cat(sf_groups) if sf_groups else torch.empty((0, n), dtype=torch.float, device=x.device) - return x_fp8, sf - - -def generate_k_grouped_contiguous(num_groups: int, m: int, n: int, major_a: MajorTypeAB, major_b: MajorTypeAB, ks_cpu: List[int], - use_ue8m0: bool = False, use_bf16: bool = False, gran_k = 128): - k = sum(ks_cpu) - grouped_layout = torch.tensor(ks_cpu, device='cuda', dtype=torch.int32) + sf = torch.cat(sf_groups) if sf_groups else torch.empty((0, mn), dtype=torch.float, device=x.device) + if major.is_mn_major(): + return x_fp8, sf - 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_cpu): - 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, grouped_layout, ks_cpu - - assert get_mk_alignment_for_contiguous_layout() % 32 == 0 - - a_fp8 = k_grouped_per_channel_cast_to_fp8(a, ks_cpu, use_ue8m0=use_ue8m0, gran_k=gran_k) - b_fp8 = k_grouped_per_channel_cast_to_fp8(b, ks_cpu, 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_cpu) * m, ), dtype=a.dtype, device=a.device) - new_b = torch.empty((sum(ks_cpu) * n, ), dtype=b.dtype, device=b.device) - prefix = 0 - for K in ks_cpu: - 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, grouped_layout, ks_cpu + data = torch.zeros((x.size(0) * mn, ), dtype=x_fp8.dtype, device=x.device) + for k, end in zip(ks_cpu, group_ends): + start = end - k + data[start * mn:end * mn] = x_fp8[start:end].T.flatten() + return data, sf.T def build_psum_layout_from_ks(real_ks: List[int], k_alignment: int) -> List[int]: - # Convert raw per-group K sizes to psum end offsets psum, prev_end = [], 0 for k in real_ks: end = align(prev_end, k_alignment) + k @@ -489,44 +548,61 @@ def build_psum_layout_from_ks(real_ks: List[int], k_alignment: int) -> List[int] return psum -def generate_k_grouped_contiguous_psum(num_groups: int, m: int, n: int, - major_a: MajorTypeAB, major_b: MajorTypeAB, - real_ks: List[int], k_alignment: int, - use_ue8m0: bool = False, use_bf16: bool = False, - gran_k: int = 128): - # psum k-grouped is SM100 MN-major only - assert (major_a, major_b) == (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor) - assert k_alignment % 32 == 0 - - # NOTES: `aligned_ks` round each group's K up to `k_alignment` (the group-start alignment), - # so the host-side `k % k_alignment == 0` check passes and matches the psum layout below. - aligned_ks = [align(k, k_alignment) for k in real_ks] - - grouped_layout = build_psum_layout_from_ks(real_ks, k_alignment) - total_k = align(grouped_layout[-1] if num_groups > 0 else 0, k_alignment) - - # Keep padded K gaps zeroed +def generate_k_grouped_contiguous(num_groups: int, m: int, n: int, major_a: MajorTypeAB, major_b: MajorTypeAB, logical_ks_cpu: List[int], + use_ue8m0: bool = False, use_bf16: bool = False, gran_k = 128, + quant_config: Optional[QuantConfig] = None, + use_psum_layout: bool = False, k_alignment: Optional[int] = None, + accumulate: bool = True, out_dtype: torch.dtype = torch.float): + assert num_groups == len(logical_ks_cpu) + k_alignment = get_mk_alignment_for_contiguous_layout() if k_alignment is None else k_alignment + host_ks_cpu = [align(k, k_alignment) for k in logical_ks_cpu] + if use_psum_layout: + assert k_alignment % 32 == 0 + logical_group_ends = build_psum_layout_from_ks(logical_ks_cpu, k_alignment) + total_k = sum(host_ks_cpu) + grouped_layout_cpu = logical_group_ends + else: + total_k = sum(host_ks_cpu) + logical_group_ends = [] + physical_start = 0 + for logical_k, host_k in zip(logical_ks_cpu, host_ks_cpu): + logical_group_ends.append(physical_start + logical_k) + physical_start += host_k + grouped_layout_cpu = host_ks_cpu + grouped_layout = torch.tensor(grouped_layout_cpu, device='cuda', dtype=torch.int32) + + # Physical psum gaps are guaranteed to be zero. A partial block must still + # stop before the next group's nonzero data. a = torch.zeros((total_k, m), device='cuda', dtype=torch.bfloat16) b = torch.zeros((total_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) + d = torch.randn((num_groups, m, n), device='cuda', dtype=out_dtype) * 32 if accumulate else \ + torch.empty((num_groups, m, n), device='cuda', dtype=out_dtype) + c = d if accumulate else None + ref_d = torch.empty_like(d) + + for i, (group_k, end) in enumerate(zip(logical_ks_cpu, logical_group_ends)): + start = end - group_k + a[start:end] = torch.randn((group_k, m), device='cuda', dtype=torch.bfloat16) + b[start:end] = torch.randn((group_k, n), device='cuda', dtype=torch.bfloat16) + ref_d[i] = (a[start:end].float().T @ b[start:end].float() + (c[i] if accumulate else 0)).to(out_dtype) - for i, k in enumerate(real_ks): - if k == 0: - ref_d[i] = c[i] - continue - start = grouped_layout[i] - k - a_g = torch.randn((k, m), device='cuda', dtype=torch.bfloat16) - b_g = torch.randn((k, n), device='cuda', dtype=torch.bfloat16) - a[start:start + k] = a_g - b[start:start + k] = b_g - ref_d[i] = c[i] + (a_g.T @ b_g) - - grouped_layout_tensor = torch.tensor(grouped_layout, device='cuda', dtype=torch.int32) if use_bf16: - return total_k, a, b, c, d, ref_d, grouped_layout_tensor, aligned_ks + assert (major_a, major_b) == (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor) + return total_k, a, b, c, d, ref_d, grouped_layout, host_ks_cpu - a_fp8 = k_grouped_per_channel_cast_to_fp8(a, real_ks, use_ue8m0=use_ue8m0, gran_k=gran_k, group_ends=grouped_layout) - b_fp8 = k_grouped_per_channel_cast_to_fp8(b, real_ks, use_ue8m0=use_ue8m0, gran_k=gran_k, group_ends=grouped_layout) - return total_k, a_fp8, b_fp8, c, d, ref_d, grouped_layout_tensor, aligned_ks + is_fp4 = quant_config is not None and (quant_config.is_fp4_a or quant_config.is_fp4_b) + if is_fp4: + assert quant_config.is_fp4_a and quant_config.is_fp4_b + assert (quant_config.gran_k_a, quant_config.gran_k_b) == (32, 32) + assert (major_a, major_b) == (MajorTypeAB.KMajor, MajorTypeAB.KMajor) + assert k_alignment % 256 == 0 + else: + assert k_alignment % 128 == 0 + + quantized_a = k_grouped_cast_fp8_fp4_with_major( + a, logical_ks_cpu if use_psum_layout else host_ks_cpu, major_a, use_ue8m0, gran_k, is_fp4, + logical_group_ends if use_psum_layout else None, host_ks_cpu) + quantized_b = k_grouped_cast_fp8_fp4_with_major( + b, logical_ks_cpu if use_psum_layout else host_ks_cpu, major_b, use_ue8m0, gran_k, is_fp4, + logical_group_ends if use_psum_layout else None, host_ks_cpu) + return total_k, quantized_a, quantized_b, c, d, ref_d, grouped_layout, host_ks_cpu diff --git a/tests/test_attention.py b/tests/test_attention.py index 1542e2b751..f7a0bb7ddf 100644 --- a/tests/test_attention.py +++ b/tests/test_attention.py @@ -67,7 +67,7 @@ def sample_mqa_cases(name: str, cases: List[tuple]) -> List[tuple]: if num_cases is None: selected = cases else: - rng = random.Random({'prefill': 0, 'paged': 100000}[name]) + rng = random.Random({'prefill': 0, 'paged': 100000, 'sparse': 200000}[name]) selected = rng.sample(cases, min(int(num_cases), len(cases))) print(f' > {name}: running {len(selected)}/{len(cases)} cases') return selected @@ -81,6 +81,15 @@ def dtype_tag(dtype: torch.dtype) -> str: return 'BF16' if dtype == torch.bfloat16 else 'FP32' +def to_mqa_weights(weights: torch.Tensor, dtype: torch.dtype) -> torch.Tensor: + element_size = torch.empty((), dtype=dtype).element_size() + stride = ceil_div(weights.size(1) * element_size, 16) * 16 // element_size + storage = torch.empty((weights.size(0), stride), device=weights.device, dtype=dtype) + result = storage[:, :weights.size(1)] + result.copy_(weights) + return result + + 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] @@ -146,7 +155,7 @@ def enumerate_mqa_logits(): for seq_len in (2048, 8192): for seq_len_kv in (8192, 65536): head_dims = (64, 128) if is_mxfp4 else (32, 64, 128) - heads = (8, 16, 32, 64) if get_arch_major() == 10 else (32, 64) + heads = (8, 12, 16, 20, 32, 64) if get_arch_major() == 10 else (32, 64) for num_heads in heads: for head_dim in head_dims: for disable_cp in (False, True): @@ -162,7 +171,7 @@ def enumerate_mqa_logits(): 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) - kernel_weights = weights.to(weights_dtype) + kernel_weights = to_mqa_weights(weights, weights_dtype) ks, ke = generate_ks_ke_tests(seq_len, seq_len_kv, disable_cp) # Calculate reference logits @@ -216,6 +225,14 @@ def enumerate_mqa_logits(): logits_again = logits_again.masked_fill(~self_mask, 0) assert_bitwise_equal(logits_again, masked_logits, 'mqa logits self-consistency') + workspace = None + if get_arch_major() == 10: + workspace = deep_gemm.get_mqa_logits_metadata(ks, ke, seq_len_kv, num_heads) + scheduled_logits = deep_gemm.fp8_fp4_mqa_logits(**kernel_kwargs, schedule_meta=workspace) + if compressed_logits: + scheduled_logits = scheduled_logits.masked_fill(~self_mask, 0) + assert_bitwise_equal(scheduled_logits, masked_logits, 'mqa logits scheduled path') + # Post process for compressed logits if compressed_logits: assert logits.size() == (seq_len, max_seqlen_k) @@ -239,17 +256,21 @@ def enumerate_mqa_logits(): # 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) * logits_dtype.itemsize + count_bytes(ks, ke) - + t = bench_kineto(lambda: deep_gemm.fp8_fp4_mqa_logits(**kernel_kwargs), 'mqa_logits') + t_scheduled = t_build = 0 + if workspace is not None: + t_scheduled = bench_kineto(lambda: deep_gemm.fp8_fp4_mqa_logits( + **kernel_kwargs, schedule_meta=workspace), 'mqa_logits') + t_build = bench_kineto(lambda: deep_gemm.get_mqa_logits_metadata( + ks, ke, seq_len_kv, num_heads), 'mqa_logits_metadata') reduce_relus = ref_cost * num_heads relu_per_sm_cycle = reduce_relus / (t * deep_gemm.get_num_sms() * 1.95 * 1e9) print(f' > Fmt={fmt:5}, Logits={dtype_tag(logits_dtype):4}, Reduce={dtype_tag(weights_dtype):4}, ' f'CMP={int(compressed_logits):1d}, SQ={seq_len:4}, SK={seq_len_kv:5}, H={num_heads:2}, D={head_dim:3}, CP={0 if disable_cp else 1}: ' - f'{tflops / t:4.0f} TFLOPS, {t * 1e6:4.0f} us, ' + f'{tflops / t:4.0f} TFLOPS, {t * 1e6:4.0f} us ' + f'(scheduled {t_scheduled * 1e6:4.0f} us, build {t_build * 1e6:4.1f} us), ' f'{(count_bytes(q_in, kv_in, kernel_weights, ks, ke) + ref_cost * logits_dtype.itemsize) / t / 1e9:4.0f} GB/s, ' - f'{relu_per_sm_cycle:4.1f} relu/cyc/SM', end='') - print(f' | clean: {clean_t * 1e6:3.0f} us, {clean_bytes / clean_t / 1e9:4.0f} GB/s' if clean_logits else '') + f'{relu_per_sm_cycle:4.1f} relu/cyc/SM') print() @@ -284,6 +305,34 @@ def ref_paged_mqa_logits(q: torch.Tensor, kv_cache: torch.Tensor, return logits +def kv_cache_cast_to_mxfp4(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + num_blocks, block_size, num_heads, head_dim = x.shape + assert num_heads == 1 and head_dim in (64, 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 kv_cache_cast_to_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + num_blocks, block_size, num_heads, head_dim = x.shape + assert num_heads == 1 and head_dim in (32, 64, 128) + x_scaled, sf = per_token_cast_to_fp8( + x.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + x_cast_back = cast_back_from_fp8( + x_scaled, sf, gran_k=32, use_packed_ue8m0=True).view(num_blocks, block_size, 1, head_dim) + + 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 test_paged_mqa_logits(): # Helper functions @@ -300,28 +349,6 @@ def kv_cache_cast_to_fp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: 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_mxfp4(x: torch.Tensor) -> torch.Tensor: - num_blocks, block_size, num_heads, head_dim = x.shape - assert num_heads == 1 and head_dim in (64, 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 kv_cache_cast_to_mxfp8(x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - num_blocks, block_size, num_heads, head_dim = x.shape - assert num_heads == 1 and head_dim in (32, 64, 128) - x_scaled, sf = per_token_cast_to_fp8(x.view(-1, head_dim), use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) - x_cast_back = cast_back_from_fp8(x_scaled, sf, gran_k=32, use_packed_ue8m0=True).view(num_blocks, block_size, 1, head_dim) - - 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 enumerate_paged_mqa_logits(): arch_major = get_arch_major() max_kv_pool_tokens = 32 * 1024 * 1024 @@ -338,7 +365,7 @@ def enumerate_paged_mqa_logits(): for batch_size in (256, 4096): for next_n in ((1, ) if is_varlen else ((1, 6) if arch_major == 10 else (1, 2))): for max_tokens_per_batch in ((6, 10) if is_varlen else (1, )): - heads = (8, 16, 32, 64) if arch_major == 10 else (32, 64) + heads = (8, 12, 16, 20, 32, 64) if arch_major == 10 else (32, 64) head_dims = (64, 128) if is_mxfp4 else ((32, 64, 128) if arch_major == 10 else (128, )) for num_heads in heads: for head_dim in head_dims: @@ -368,7 +395,7 @@ def enumerate_paged_mqa_logits(): # Generate random inputs q = torch.randn((batch_size, next_n, num_heads, head_dim), device='cuda', dtype=torch.bfloat16) weights = torch.randn((batch_size * next_n, num_heads), device='cuda', dtype=torch.float) - kernel_weights = weights.to(weights_dtype) + kernel_weights = to_mqa_weights(weights, weights_dtype) 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: @@ -444,10 +471,14 @@ def enumerate_paged_mqa_logits(): assert block_table.min().item() >= 0 assert block_table.max().item() < num_total_blocks assert context_lens_nextn.max().item() <= max_model_len + metadata_kwargs = dict( + context_lens=context_lens_nextn, block_kv=block_kv, + num_sms=deep_gemm.get_num_sms(), indices=indices, + ) kernel_kwargs = dict( q=q_in, kv_cache=kv_in, weights=kernel_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), + schedule_meta=deep_gemm.get_paged_mqa_logits_metadata(**metadata_kwargs), max_context_len=max_model_len, clean_logits=clean_logits, logits_dtype=logits_dtype, indices=indices, ) @@ -482,16 +513,20 @@ def enumerate_paged_mqa_logits(): kv_sum_lens = seq_sum_lens if is_varlen else sum_lens total_bytes = q_weight_bytes + 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')) + metadata_t = bench_kineto( + lambda: deep_gemm.get_paged_mqa_logits_metadata(**metadata_kwargs), + 'paged_mqa_logits_metadata', + ) + t = bench_kineto(lambda: deep_gemm.fp8_fp4_paged_mqa_logits(**kernel_kwargs), 'paged_mqa_logits') reduce_relus = sum_lens * next_n * num_heads relu_per_sm_cycle = reduce_relus / (t * deep_gemm.get_num_sms() * 1.95 * 1e9) next_n_desc = f'MaxTPR={max_tokens_per_batch:2}' if is_varlen else f'NextN ={raw_next_n:2}' print(f' > Fmt={fmt:5}, Logits={dtype_tag(logits_dtype):4}, Reduce={dtype_tag(weights_dtype):4}, ' f'VAR={int(is_varlen):1d}, PAGE_KV={block_kv:2}, BSZ={raw_batch_size:4}, {next_n_desc}, H={num_heads:2}, D={head_dim:3}, L={avg_kv:5}: ' - f'{tflops_calc / t:4.0f} TFLOPS, {t * 1e6:4.0f} us, {total_bytes / t / 1e9:4.0f} GB/s, {relu_per_sm_cycle:4.1f} relu/cyc/SM', end='') - print(f' | clean: {clean_t*1e6:3.0f} us' if clean_logits else '') + f'{tflops_calc / t:4.0f} TFLOPS, {t * 1e6:4.0f} us, Metadata={metadata_t * 1e6:4.0f} us, ' + f'{total_bytes / t / 1e9:4.0f} GB/s, {relu_per_sm_cycle:4.1f} relu/cyc/SM') - del kernel_kwargs, logits, ref_neginf_mask, positions + del metadata_kwargs, kernel_kwargs, logits, ref_neginf_mask, positions del q_in, q_simulated, kv_in, kv_simulated, weights, kernel_weights, context_lens, context_lens_nextn, block_table if is_mxfp4 or is_mxfp8: del q_q @@ -501,6 +536,274 @@ def enumerate_paged_mqa_logits(): print() +def make_sparse_kv_block_indices(context_lens: List[int], request_indices: List[int] | None, + sparse_block_kv: int, num_max_sparse_blocks: int, + seed: int, context_starts: List[int] | None = None) -> Tuple[torch.Tensor, List[int]]: + rng = random.Random(seed) + request_indices = [0] * len(context_lens) if request_indices is None else request_indices + context_starts = [0] * len(context_lens) if context_starts is None else context_starts + indices, num_blocks_per_q = [], [] + previous_blocks, previous_request_idx = None, None + for context_start, context_len, request_idx in zip(context_starts, context_lens, request_indices): + first_block = context_start // sparse_block_kv + num_available_blocks = ceil_div(max(0, context_len - context_start), sparse_block_kv) + block_end = first_block + num_available_blocks + num_sparse_blocks = min(num_available_blocks, num_max_sparse_blocks) + if num_sparse_blocks == num_available_blocks: + blocks = list(range(first_block, block_end)) + elif request_idx != previous_request_idx: + blocks = rng.sample(range(first_block, block_end), num_sparse_blocks) + else: + previous = [block_idx for block_idx in previous_blocks if first_block <= block_idx < block_end] + retained = rng.sample(previous, min(round(num_sparse_blocks * 0.8), len(previous))) + retained_set = set(retained) + replacements = set() + while len(retained) + len(replacements) < num_sparse_blocks: + block_idx = rng.randrange(first_block, block_end) + if block_idx not in retained_set: + replacements.add(block_idx) + blocks = retained + list(replacements) + blocks.sort() + indices.append(blocks + [blocks[-1] if blocks else 0] * (num_max_sparse_blocks - num_sparse_blocks)) + num_blocks_per_q.append(num_sparse_blocks) + previous_blocks, previous_request_idx = blocks, request_idx + return torch.tensor(indices, device='cuda', dtype=torch.int32), num_blocks_per_q + + +@test_filter(lambda: get_arch_major() == 10) +def test_sparse_mqa_logits() -> None: + num_heads, head_dim = 32, 128 + page_kv = 64 + + def enumerate_sparse_mqa_logits(): + avg_kv_lens = (4 * 1024, 8 * 1024, 16 * 1024, 32 * 1024, 64 * 1024, 128 * 1024, 256 * 1024) + num_sms = deep_gemm.get_num_sms() + + for fmt in ('mxfp4', 'mxfp8'): + # Contiguous KV + for num_max_sparse_blocks in (2048, 1024, 512): + for num_q_tokens in (8192, ): + for avg_kv_len in avg_kv_lens: + for use_unaligned_ks in (False, True): + yield fmt, False, num_q_tokens, avg_kv_len, 8, num_max_sparse_blocks, use_unaligned_ks + + # Paged varlen KV + for num_max_sparse_blocks in (2048, 1024, 512): + for num_q_tokens in (512, ): + for avg_kv_len in avg_kv_lens: + yield fmt, True, num_q_tokens, avg_kv_len, 8, num_max_sparse_blocks, False + + # Small Q counts and both sparse block sizes. + split_kv = 640 if fmt == 'mxfp4' else 512 + for sparse_block_kv in (8, 16): + for use_unaligned_ks in (False, True): + yield fmt, False, 9, split_kv, sparse_block_kv, 128, use_unaligned_ks + + # Contiguous metadata edge cases use MXFP4 + for case in ( + (1, 1, 16, 4), (2, 639, 16, 1024), (2, 0, 16, 4), + (3, 640, 16, 1024), (5, 641, 16, 1024), + (num_sms - 1, 16 * 1024 + 3, 16, 1024), + (num_sms, 16 * 1024 + 3, 16, 1024), + (num_sms + 1, 16 * 1024 + 3, 16, 1024), + (2 * num_sms + 1, 64 * 1024 + 7, 16, 4096), + ): + for use_unaligned_ks in (False, True): + yield 'mxfp4', False, *case, use_unaligned_ks + + # Paged metadata edge cases use MXFP4 + for case in ( + (True, 1, 64, 16, 8), (True, 3, 64, 16, 8), + (True, num_sms - 1, 16 * 1024 + 3, 16, 1024), + (True, num_sms + 1, 16 * 1024 + 3, 16, 1024), + (True, 2 * num_sms + 1, 16 * 1024 + 3, 16, 1024), + (True, 3, 1024 * 1024, 8, 2048), + ): + yield 'mxfp4', *case, False + for use_unaligned_ks in (False, True): + yield 'mxfp8', False, 3, 640, 16, 1024, use_unaligned_ks + yield 'mxfp8', True, 3, 64, 16, 8, False + + print('Testing MXFP4/MXFP8 Sparse MQA Logits:') + torch.manual_seed(0) + cases = sample_mqa_cases('sparse', list(enumerate_sparse_mqa_logits())) + aligned_sparse_times = {} + for fmt, is_paged, num_q_tokens, avg_kv_len, sparse_block_kv, num_max_sparse_blocks, use_unaligned_ks in cases: + is_mxfp4 = fmt == 'mxfp4' + cast_fwd = per_token_cast_to_fp4 if is_mxfp4 else per_token_cast_to_fp8 + kv_cache_cast = kv_cache_cast_to_mxfp4 if is_mxfp4 else kv_cache_cast_to_mxfp8 + elem_dim = head_dim // 2 if is_mxfp4 else head_dim + rng = random.Random(num_q_tokens * 1000000 + avg_kv_len + sparse_block_kv) + request_sizes = [] + if is_paged: + remaining_q_tokens = num_q_tokens + while remaining_q_tokens > 0: + request_size = min(rng.randint(2, 6), remaining_q_tokens) + request_sizes.append(request_size) + remaining_q_tokens -= request_size + else: + num_requests = min(rng.randint(2, 4), num_q_tokens) + request_ends = sorted(rng.sample(range(1, num_q_tokens), num_requests - 1)) + [num_q_tokens] + request_sizes = [request_end - request_begin + for request_begin, request_end in zip([0] + request_ends, request_ends)] + request_indices = [request_idx for request_idx, request_size in enumerate(request_sizes) + for _ in range(request_size)] + if is_paged: + context_starts = [0] * num_q_tokens + batch_size = len(request_sizes) + request_context_lens = [rng.randint(int(0.7 * avg_kv_len) // sparse_block_kv, + int(1.3 * avg_kv_len) // sparse_block_kv) * sparse_block_kv + for _ in range(batch_size)] + context_lens = [context_len + q_offset + for context_len, request_size in zip(request_context_lens, request_sizes) + for q_offset in range(request_size)] + else: + aligned_starts, unaligned_starts, context_lengths = [], [], [] + aligned_kv_end = unaligned_kv_end = 0 + for request_size in request_sizes: + aligned_start = ceil_div(aligned_kv_end, sparse_block_kv) * sparse_block_kv + unaligned_start = ceil_div(unaligned_kv_end, sparse_block_kv) * sparse_block_kv + rng.randrange(1, sparse_block_kv) + aligned_starts.extend([aligned_start] * request_size) + unaligned_starts.extend([unaligned_start] * request_size) + context_lengths.extend(avg_kv_len + q_offset for q_offset in range(request_size)) + aligned_kv_end = aligned_start + context_lengths[-1] + unaligned_kv_end = unaligned_start + context_lengths[-1] + assert all(start % sparse_block_kv == 0 for start in aligned_starts) + assert all(start % sparse_block_kv != 0 for start in unaligned_starts) + context_starts = unaligned_starts if use_unaligned_ks else aligned_starts + context_lens = [start + length for start, length in zip(context_starts, context_lengths)] + num_kv_tokens = max(1, aligned_kv_end, unaligned_kv_end) + + q_shape = (num_q_tokens, 1, num_heads) if is_paged else (num_q_tokens, num_heads) + q_fp, q_sf = cast_fwd( + torch.randn((*q_shape, head_dim), device='cuda', dtype=torch.bfloat16).view(-1, head_dim), + use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + q = q_fp.view(*q_shape, elem_dim), q_sf.view(*q_shape) + weights = to_mqa_weights(torch.randn((num_q_tokens, num_heads), device='cuda', dtype=torch.bfloat16), + torch.bfloat16) + sparse_indices, num_sparse_blocks = make_sparse_kv_block_indices( + context_lens, request_indices, sparse_block_kv, num_max_sparse_blocks, avg_kv_len, context_starts) + + if is_paged: + max_context_lens = [context_len + request_size - 1 + for context_len, request_size in zip(request_context_lens, request_sizes)] + num_pages_per_request = [ceil_div(context_len, page_kv) for context_len in max_context_lens] + max_num_pages, num_pages = max(num_pages_per_request), sum(num_pages_per_request) + page_bytes = page_kv * (elem_dim + 4) + page_stride_bytes = ceil_div(page_bytes, 512) * 512 + kv_storage = torch.empty((num_pages, page_stride_bytes), device='cuda', dtype=torch.uint8) + kv_cache = kv_storage.as_strided((num_pages, page_kv, 1, elem_dim + 4), + (page_stride_bytes, elem_dim + 4, elem_dim + 4, 1)) + for page_begin in range(0, num_pages, 16 * 1024): + num_pages_to_copy = min(16 * 1024, num_pages - page_begin) + kv_pages, kv_pages_reference = kv_cache_cast(torch.randn( + (num_pages_to_copy, page_kv, 1, head_dim), device='cuda', dtype=torch.bfloat16)) + kv_cache[page_begin:page_begin + num_pages_to_copy].copy_(kv_pages) + del kv_pages, kv_pages_reference + + indices = torch.tensor(request_indices, device='cuda', dtype=torch.int32) + context_lens_tensor = torch.tensor(context_lens, device='cuda', dtype=torch.int32) + page_pool = torch.randperm(num_pages, device='cuda', dtype=torch.int32) + request_block_table = torch.zeros((batch_size, max_num_pages), device='cuda', dtype=torch.int32) + page_begin = 0 + for request_idx, num_request_pages in enumerate(num_pages_per_request): + page_end = page_begin + num_request_pages + request_block_table[request_idx, :num_request_pages] = page_pool[page_begin:page_end] + page_begin = page_end + block_table = request_block_table[indices.long()].contiguous() + metadata = deep_gemm.get_paged_sparse_mqa_logits_metadata( + context_lens_tensor, block_table, indices, page_kv, sparse_indices, q[0].dtype, sparse_block_kv) + sparse_kwargs = dict(q=q, kv_cache=kv_cache, weights=weights, metadata=metadata, + num_max_sparse_blocks=num_max_sparse_blocks, sparse_block_kv=sparse_block_kv) + context_lens_2d = context_lens_tensor.view(-1, 1) + full_kwargs = dict( + q=q, kv_cache=kv_cache, weights=weights, context_lens=context_lens_2d, block_table=block_table, + schedule_meta=deep_gemm.get_paged_mqa_logits_metadata( + context_lens_2d, page_kv, deep_gemm.get_num_sms(), indices), + max_context_len=max(context_lens), clean_logits=False, + logits_dtype=torch.bfloat16, indices=indices) + run_sparse = lambda: deep_gemm.fp8_fp4_paged_sparse_mqa_logits(**sparse_kwargs) + run_full = lambda: deep_gemm.fp8_fp4_paged_mqa_logits(**full_kwargs) + sparse_kernel_name, full_kernel_name = 'sm100_paged_sparse_mqa_logits', 'paged_mqa_logits' + case = f'Paged BSZ={batch_size:3}, SQ={num_q_tokens:4}' + else: + kv_fp, kv_sf = cast_fwd( + torch.randn((num_kv_tokens, head_dim), device='cuda', dtype=torch.bfloat16), + use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + kv = kv_fp, kv_sf.view(num_kv_tokens) + starts = torch.tensor(context_starts, device='cuda', dtype=torch.int32) + ends = torch.tensor(context_lens, device='cuda', dtype=torch.int32) + metadata_kwargs = dict( + cu_seq_len_k_start=starts, cu_seq_len_k_end=ends, num_kv_tokens=num_kv_tokens, + sparse_kv_block_indices=sparse_indices, qk_dtype=q[0].dtype, sparse_block_kv=sparse_block_kv) + if use_unaligned_ks: + metadata_kwargs['use_unaligned_ks'] = True + metadata = deep_gemm.get_sparse_mqa_logits_metadata(**metadata_kwargs) + sparse_kwargs = dict(q=q, kv=kv, weights=weights, metadata=metadata, + num_max_sparse_blocks=num_max_sparse_blocks, sparse_block_kv=sparse_block_kv) + if use_unaligned_ks: + sparse_kwargs['use_unaligned_ks'] = True + full_kwargs = dict(q=q, kv=kv, weights=weights, cu_seq_len_k_start=starts, cu_seq_len_k_end=ends, + clean_logits=False, max_seqlen_k=num_kv_tokens, logits_dtype=torch.bfloat16) + run_sparse = lambda: deep_gemm.fp8_fp4_sparse_mqa_logits(**sparse_kwargs) + run_full = lambda: deep_gemm.fp8_fp4_mqa_logits(**full_kwargs) + sparse_kernel_name, full_kernel_name = 'sm100_sparse_mqa_logits', 'mqa_logits' + case = f'Contiguous SQ={num_q_tokens:4}, UnalignedKS={int(use_unaligned_ks)}' + + sparse_logits, full_logits = run_sparse(), run_full() + sparse_output_bytes = count_bytes(sparse_logits) + kv_offsets = torch.arange(sparse_block_kv, device='cuda', dtype=torch.int32) + context_starts_tensor = torch.tensor(context_starts, device='cuda', dtype=torch.int32) + block_offsets = context_starts_tensor % sparse_block_kv + token_indices = (sparse_indices.unsqueeze(-1) * sparse_block_kv + + block_offsets[:, None, None] + kv_offsets).flatten(1).long() + num_sparse_blocks = torch.tensor(num_sparse_blocks, device='cuda') + num_sparse_tokens = num_sparse_blocks * sparse_block_kv + valid_mask = torch.arange(token_indices.size(1), device='cuda')[None, :] < num_sparse_tokens[:, None] + valid_mask &= token_indices >= context_starts_tensor[:, None] + valid_mask &= token_indices < torch.tensor(context_lens, device='cuda')[:, None] + sparse_logits = sparse_logits[valid_mask] + full_token_indices = token_indices - context_starts_tensor[:, None] + full_logits = full_logits.gather( + 1, full_token_indices.clamp(min=0, max=full_logits.size(1) - 1))[valid_mask] + assert_bitwise_equal(sparse_logits, full_logits, 'sparse MQA logits') + + for _ in range(30): + assert_bitwise_equal(run_sparse()[valid_mask], sparse_logits, 'sparse MQA logits self-consistency') + + sparse_t = bench_kineto(run_sparse, sparse_kernel_name) + full_t = bench_kineto(run_full, full_kernel_name) + comparison = '' + if not is_paged: + key = fmt, num_q_tokens, avg_kv_len, sparse_block_kv, num_max_sparse_blocks + if use_unaligned_ks and key in aligned_sparse_times: + comparison = f', unaligned/aligned {sparse_t / aligned_sparse_times[key]:.2f}x' + elif not use_unaligned_ks: + aligned_sparse_times[key] = sparse_t + valid_block_mask = torch.arange(num_max_sparse_blocks, device='cuda')[None, :] < num_sparse_blocks[:, None] + if is_paged: + request_indices_2d = indices[:, None].expand_as(sparse_indices) + block_keys = torch.stack((request_indices_2d[valid_block_mask], sparse_indices[valid_block_mask]), dim=-1) + num_union_blocks = torch.unique(block_keys, dim=0).size(0) + else: + block_starts = sparse_indices * sparse_block_kv + block_offsets[:, None] + num_union_blocks = block_starts[valid_block_mask].unique().numel() + num_sum_blocks = num_sparse_blocks.sum().item() + kv_bytes_per_block = sparse_block_kv * (elem_dim + 4) + total_bytes = count_bytes(q, weights, metadata) + sparse_output_bytes + num_union_blocks * kv_bytes_per_block + tflops = 2 * num_sum_blocks * sparse_block_kv * num_heads * head_dim / 1e12 + reduce_relus = num_sum_blocks * sparse_block_kv * num_heads + relu_per_sm_cycle = reduce_relus / (sparse_t * deep_gemm.get_num_sms() * 1.95 * 1e9) + print(f' > Fmt={fmt:5}, {case}, KV={avg_kv_len:7}, SPARSE_BLOCK_KV={sparse_block_kv:2}, ' + f'MAX_BLOCKS={num_max_sparse_blocks:4}: sparse {sparse_t * 1e6:5.1f} us, ' + f'{tflops / sparse_t:4.0f} TFLOPS, ' + f'{total_bytes / sparse_t / 1e9:4.0f} GB/s, ' + f'{relu_per_sm_cycle:4.1f} relu/cyc/SM ', + f'(full {full_t * 1e6:6.1f} us, {full_t / sparse_t:5.2f}x{comparison})') + torch.cuda.empty_cache() + print() + + if __name__ == '__main__': @@ -510,3 +813,4 @@ def enumerate_paged_mqa_logits(): test_gemm_skip_head_mid() test_mqa_logits() test_paged_mqa_logits() + test_sparse_mqa_logits() diff --git a/tests/test_bf16.py b/tests/test_bf16.py index eec58bee45..ccb62f876e 100644 --- a/tests/test_bf16.py +++ b/tests/test_bf16.py @@ -8,42 +8,47 @@ calc_diff, count_bytes ) from utils import ( + assert_direct_output_matches_fp32_accumulation, assert_psum_zero_padding, ) 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, + get_arch_major, + enumerate_normal, enumerate_batched_syrk_symm, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous, enumerate_k_grouped_contiguous_test_variants, generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous, - generate_k_grouped_contiguous_psum, - get_mk_alignment_for_contiguous_layout ) def test_gemm() -> None: print('Testing GEMM:') scores = [] + use_alpha_options = (False, True) if get_arch_major() == 10 else (False,) for kernel_type, _, m, n, k, major_a, major_b, accumulate, out_dtype in enumerate_normal(torch.bfloat16): + deep_gemm.use_deterministic_algorithms(True) 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}') + for use_alpha in use_alpha_options: + alpha = random.uniform(-1.0, 1.0) if use_alpha else None + a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, + kernel_type, use_bf16=True, alpha=alpha) + 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, alpha=alpha) + diff = calc_diff(d, ref_d) + assert diff < 1e-5, (f'{m=}, {n=}, {k=}, {major_opt=}, {accumulate=}, {out_dtype=}, ' + f'{use_alpha=}, {alpha=}, {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=c), ('nvjet', 'reduce'), suppress_kineto_output=True) + deep_gemm.use_deterministic_algorithms(False) + cublas_t, split_k_t = bench_kineto(lambda: deep_gemm.bf16_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 | ' @@ -66,8 +71,8 @@ def test_m_grouped_gemm_contiguous() -> None: 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) + m, a, b, grouped_layout, d, ref_d, valid_mask = 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() @@ -75,20 +80,12 @@ def test_m_grouped_gemm_contiguous() -> None: 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, ensure_zero_padding=ensure_zero_padding) - 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}, ' - f'alias={test_alias}, {ensure_zero_padding=}') - if ensure_zero_padding: - assert_psum_zero_padding(a, d, grouped_layout, 'BF16') - 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) + diff = calc_diff(d[valid_mask], ref_d[valid_mask]) + assert diff < 1e-5, f'{m=}, {n=}, {k=}, {major_opt}, {diff:.5f}, alias={test_alias}, {ensure_zero_padding=}' + if use_psum_layout and ensure_zero_padding: + assert_psum_zero_padding(a, d, grouped_layout, 'BF16') + m, a, b, grouped_layout, d, ref_d, valid_mask = 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(): @@ -118,40 +115,30 @@ def test_m_grouped_gemm_masked() -> None: 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) + a, b, grouped_layout, d, ref_d, valid_mask = generate_m_grouped_masked( + num_groups, max_m, expected_m_per_group, n, k, + use_bf16=True, use_psum_layout=use_psum_layout) - # noinspection PyShadowingNames def test_func(): if use_psum_layout: - deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a_psum, b, d_psum, psum_m, + deep_gemm.m_grouped_bf16_gemm_nt_contiguous(a, b, d, grouped_layout, 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) + deep_gemm.m_grouped_bf16_gemm_nt_masked(a, b, d, grouped_layout, 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}' + diff = calc_diff(d[valid_mask], ref_d[valid_mask]) + assert diff < 1e-5, f'{max_m=}, {n=}, {k=}, {num_groups=}, {diff:.5f}' # Test performance with fixed shapes - valid_m = masked_m.sum().item() + valid_m = int(valid_mask.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) + 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}): ' @@ -164,47 +151,54 @@ def test_func(): def test_k_grouped_gemm_contiguous() -> None: print('Testing k-grouped contiguous GEMM:') - for num_groups, m, n, major_a, major_b, real_ks_cpu, aligned_ks_cpu, _, _, alignment, use_psum_layout in enumerate_k_grouped_contiguous(torch.bfloat16): - include_k_tail = get_arch_major() == 10 and alignment == 32 - for test_real_ks_cpu, test_aligned_ks_cpu, _, test_k_tail in enumerate_k_grouped_contiguous_test_variants(real_ks_cpu, alignment, use_psum_layout, include_k_tail): - if use_psum_layout: - total_k, a, b, c, d, ref_d, grouped_layout, _ = generate_k_grouped_contiguous_psum(num_groups, m, n, major_a, major_b, test_real_ks_cpu, k_alignment=alignment, use_bf16=True, gran_k=alignment) - else: - total_k, a, b, c, d, ref_d, grouped_layout, _ = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, test_aligned_ks_cpu, use_bf16=True) - c_orig = c.clone() if use_psum_layout else None - deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, test_aligned_ks_cpu, grouped_layout, c, use_psum_layout=use_psum_layout) - - diff = calc_diff(d, ref_d) - assert diff < 1e-5, f'{m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, {test_aligned_ks_cpu=}, {use_psum_layout=}, {test_k_tail=}, {diff:.7f}' - - # Unsynced psum paths - if use_psum_layout: - c.copy_(c_orig) - deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, None, grouped_layout, c, - use_psum_layout=True) - diff = calc_diff(d, ref_d) - assert diff < 1e-5, f'None ks_cpu path: {m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, {diff:.7f}' - - c.copy_(c_orig) - deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, [], grouped_layout, c, - use_psum_layout=True) - diff = calc_diff(d, ref_d) - assert diff < 1e-5, f'empty ks_cpu path: {m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, {diff:.7f}' + for num_groups, m, n, major_a, major_b, real_ks_cpu, _, _, _, alignment, use_psum_layout, accumulate, out_dtype in \ + enumerate_k_grouped_contiguous(torch.bfloat16): + for test_real_ks_cpu in enumerate_k_grouped_contiguous_test_variants(real_ks_cpu): + total_k, a, b, c, d, ref_d, grouped_layout, host_ks_cpu = generate_k_grouped_contiguous( + num_groups, m, n, major_a, major_b, test_real_ks_cpu, use_bf16=True, + use_psum_layout=use_psum_layout, k_alignment=alignment, + accumulate=accumulate, out_dtype=out_dtype) + initial_d = d.clone() + if not accumulate: + initial_d.fill_(float('nan')) + host_ks_options = (host_ks_cpu, None, []) if use_psum_layout else (host_ks_cpu, ) + for test_host_ks_cpu in host_ks_options: + d.copy_(initial_d) + deep_gemm.k_grouped_bf16_gemm_tn_contiguous( + a, b, d, test_host_ks_cpu, grouped_layout, c, use_psum_layout=use_psum_layout) + if accumulate: + diff = calc_diff(d, ref_d) + assert diff < 1e-5, (f'{m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, ' + f'{test_host_ks_cpu=}, {use_psum_layout=}, {accumulate=}, ' + f'{out_dtype=}, {diff:.7f}') + else: + case_label = (f'BF16 K-grouped direct output, {m=}, {n=}, {total_k=}, ' + f'{test_real_ks_cpu=}, {test_host_ks_cpu=}, {use_psum_layout=}, ' + f'{out_dtype=}') + assert_direct_output_matches_fp32_accumulation( + d, + lambda output, accumulator: deep_gemm.k_grouped_bf16_gemm_tn_contiguous( + a, b, output, test_host_ks_cpu, grouped_layout, accumulator, + use_psum_layout=use_psum_layout), + case_label) # Test performance - if use_psum_layout: - total_k, a, b, c, d, ref_d, grouped_layout, _ = generate_k_grouped_contiguous_psum(num_groups, m, n, major_a, major_b, real_ks_cpu, k_alignment=alignment, use_bf16=True, gran_k=alignment) - else: - total_k, a, b, c, d, ref_d, grouped_layout, _ = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, aligned_ks_cpu, use_bf16=True) + _, a, b, c, d, _, grouped_layout, host_ks_cpu = generate_k_grouped_contiguous( + num_groups, m, n, major_a, major_b, real_ks_cpu, use_bf16=True, + use_psum_layout=use_psum_layout, k_alignment=alignment, + accumulate=accumulate, out_dtype=out_dtype) # noinspection PyShadowingNames def test_func(): - deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, aligned_ks_cpu, grouped_layout, c, use_psum_layout=use_psum_layout) + deep_gemm.k_grouped_bf16_gemm_tn_contiguous(a, b, d, host_ks_cpu, grouped_layout, c, use_psum_layout=use_psum_layout) t = bench_kineto(test_func, 'bf16_gemm', suppress_kineto_output=True) - print(f' > Perf ({num_groups=:2}, m={m:5}, n={n:5}, k={total_k:5}, align={alignment:3}, psum={int(use_psum_layout)}): ' + logical_k = sum(real_ks_cpu) + out_opt = 'FP32' if out_dtype == torch.float else 'BF16' + print(f' > Perf ({num_groups=:2}, m={m:5}, n={n:5}, k={logical_k:5}, align={alignment:3}, ' + f'psum={int(use_psum_layout)}, acc={int(accumulate)}, {out_opt}): ' f'{t * 1e6:4.0f} us | ' - f'{2 * m * n * total_k / t / 1e12:4.0f} TFLOPS | ' + f'{2 * m * n * logical_k / t / 1e12:4.0f} TFLOPS | ' f'{count_bytes(a, b, c, d) / 1e9 / t:4.0f} GB/s') print() @@ -212,18 +206,24 @@ def test_func(): def test_cublaslt_gemm() -> None: print('Testing cuBLASLt GEMM:') + use_alpha_options = (False, True) 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=c) - diff = calc_diff(d, ref_d) # 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=})' + for use_alpha in use_alpha_options: + alpha = random.uniform(-1.0, 1.0) if use_alpha else None + a, b, c, d, ref_d = generate_normal(m, n, k, major_a, major_b, accumulate, out_dtype, + kernel_type, use_bf16=True, alpha=alpha) + deep_gemm.use_deterministic_algorithms(False) + deep_gemm.bf16_gemm_nt(a, b, d, c=c, alpha=alpha) + diff = calc_diff(d, ref_d) + assert diff < threshold, (f'{diff=}, {use_alpha=}, {alpha=}, ' + f'({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 = t_nvjet + t_gemv + t_gemm @@ -234,6 +234,85 @@ def test_cublaslt_gemm() -> None: print() +def test_cublaslt_batched_syrk() -> None: + print('Testing cuBLASLt batched SYRK:') + for num_batches, m, k, out_dtype in enumerate_batched_syrk_symm(): + out_opt = 'FP32' if out_dtype == torch.float else 'BF16' + rows, cols = min(m, k), max(m, k) + threshold = 6e-7 + for k_major in (True, False): + for batch_shape, padding in (((), (3, 5)), ((num_batches,), (3, 5)), ((num_batches,), (0, 0))): + shape = (rows, cols) if k_major else (cols, rows) + a = torch.randn(*batch_shape, shape[0] + padding[0], shape[1] + padding[1], + device='cuda', dtype=out_dtype)[..., :shape[0], :shape[1]] + if not k_major: + a = a.mT + major_opt = 'N' if a.stride(-1) == 1 else 'T' + + d = torch.empty(*batch_shape, rows + padding[0], rows + padding[1], + device='cuda', dtype=out_dtype)[..., :rows, :rows] + deep_gemm.batched_syrk(a, d) + ref_d = (a.float() @ a.float().mT).to(out_dtype) + diff = calc_diff(d, ref_d) + assert diff < threshold, ( + f'{diff=}, ({batch_shape=}, {padding=}, {rows=}, {cols=}, {major_opt=}, {out_dtype=})' + ) + + t_nvjet, t_gemv, t_gemm = bench_kineto( + lambda: deep_gemm.batched_syrk(a, d), + ('nvjet', 'gemv', 'gemm'), suppress_kineto_output=True) + t = t_nvjet + t_gemv + t_gemm + print(f' > Perf (batch={num_batches}, m={rows:6}, n={rows:6}, k={cols:6}, ' + f'layout={major_opt}, {out_opt}): ' + f'{t * 1e6:5.0f} us | ' + f'{2 * num_batches * rows * rows * cols / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, d) / 1e9 / t:4.0f} GB/s') + print() + + +def test_cublaslt_batched_symm() -> None: + print('Testing cuBLASLt batched SYMM:') + for num_batches, m, k, out_dtype in enumerate_batched_syrk_symm(): + out_opt = 'FP32' if out_dtype == torch.float else 'BF16' + rows, cols = min(m, k), max(m, k) + threshold = 6e-7 + for k_major_a, k_major_b in ((True, True), (True, False), (False, True), (False, False)): + for batch_shape, padding in (((), (3, 5)), ((num_batches,), (3, 5)), ((num_batches,), (0, 0))): + a = torch.randn(*batch_shape, rows + padding[0], rows + padding[1], + device='cuda', dtype=out_dtype)[..., :rows, :rows] + if not k_major_a: + a = a.mT + a.copy_(a + a.mT) + + shape_b = (rows, cols) if k_major_b else (cols, rows) + b = torch.randn(*batch_shape, shape_b[0] + padding[0], shape_b[1] + padding[1], + device='cuda', dtype=out_dtype)[..., :shape_b[0], :shape_b[1]] + if not k_major_b: + b = b.mT + major_opt = 'N' if a.stride(-1) == 1 else 'T' + major_opt += 'N' if b.stride(-1) == 1 else 'T' + + d = torch.empty(*batch_shape, rows + padding[0], cols + padding[1], + device='cuda', dtype=out_dtype)[..., :rows, :cols] + deep_gemm.batched_symm(a, b, d) + ref_d = (a.float() @ b.float()).to(out_dtype) + diff = calc_diff(d, ref_d) + assert diff < threshold, ( + f'{diff=}, ({batch_shape=}, {padding=}, {rows=}, {cols=}, {major_opt=}, {out_dtype=})' + ) + + t_nvjet, t_gemv, t_gemm = bench_kineto( + lambda: deep_gemm.batched_symm(a, b, d), + ('nvjet', 'gemv', 'gemm'), suppress_kineto_output=True) + t = t_nvjet + t_gemv + t_gemm + print(f' > Perf (batch={num_batches}, m={rows:6}, n={cols:6}, k={rows:6}, ' + f'layout={major_opt}, {out_opt}): ' + f'{t * 1e6:5.0f} us | ' + f'{2 * num_batches * rows * rows * cols / 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) @@ -248,3 +327,5 @@ def test_cublaslt_gemm() -> None: test_k_grouped_gemm_contiguous() test_cublaslt_gemm() + test_cublaslt_batched_syrk() + test_cublaslt_batched_symm() diff --git a/tests/test_einsum.py b/tests/test_einsum.py index 57f5459288..1a47724a4e 100644 --- a/tests/test_einsum.py +++ b/tests/test_einsum.py @@ -8,7 +8,7 @@ get_arch_major, test_filter ) from deep_gemm.utils.math import ( - ceil_div, + align, cast_back_from_fp8, ceil_div, per_block_cast_to_fp8, per_channel_cast_to_fp8, per_token_cast_to_fp8 ) @@ -44,12 +44,16 @@ def test_bhr_hdr_bhd(): 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) + times = [] + for use_cublaslt in (False, True): + deep_gemm.use_deterministic_algorithms(not use_cublaslt) + 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 + kernel_name = 'nvjet' if use_cublaslt else 'gemm' + times.append(bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z), kernel_name, suppress_kineto_output=True)) + t, t_cublaslt = times 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 | ' @@ -66,12 +70,17 @@ def test_bhd_hdr_bhr(): 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) + times = [] + for use_cublaslt in (False, True): + deep_gemm.use_deterministic_algorithms(not use_cublaslt) + 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 + kernel_name = 'nvjet' if use_cublaslt else 'gemm' + times.append(bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z), kernel_name, suppress_kineto_output=True)) + + t, t_cublaslt = times 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 | ' @@ -80,8 +89,30 @@ def test_bhd_hdr_bhr(): print() +def test_bhd_bhr_hdr(): + print('Testing "bhd, bhr -> hdr":') + 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) + 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.float(), y.float()) + + z = z_0.clone() + deep_gemm.einsum('bhd,bhr->hdr', x, y, z, z) + assert calc_diff(z, ref_z) < 1e-10 + + t = bench_kineto(lambda: deep_gemm.einsum('bhd,bhr->hdr', x, y, z, z), '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, z)) / t / 1e9:4.0f} GB/s') + print() + + def test_fp8_bhr_hdr_bhd(use_ue8m0: bool = True): print('Testing FP8 "bhr, hdr -> bhd":') + deep_gemm.use_deterministic_algorithms(False) 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) @@ -95,12 +126,36 @@ def test_fp8_bhr_hdr_bhd(use_ue8m0: bool = True): 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 + output_dtypes = (torch.bfloat16, torch.float8_e4m3fn) if get_arch_major() == 10 else (torch.bfloat16,) + for output_dtype in output_dtypes: + if output_dtype == torch.float8_e4m3fn: + output = torch.empty_like(z, dtype=output_dtype) + sfd = torch.empty_strided( + (b, ceil_div(h * d, 32 * 4)), (1, align(b, 4)), dtype=torch.int32, device='cuda') + deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, (output, sfd)) + restored_output = cast_back_from_fp8( + output.view(b, h * d), sfd.contiguous(), gran_k=32, use_packed_ue8m0=True) + diff = calc_diff(restored_output, ref_z.reshape(b, h * d)) + + # The fused output must align with the production pipeline (a BF16 D + # followed by a per-token cast): the epilogue rounds the accumulator into + # BF16 before the same amax/scale/rounding steps, so both the FP8 data + # and the packed SF words must match bitwise + d_bf16 = torch.empty((b, h, d), device='cuda', dtype=torch.bfloat16) + deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, d_bf16) + ref_fp8, ref_sfd = per_token_cast_to_fp8(d_bf16.view(b, h * d), use_ue8m0=True, + gran_k=32, use_packed_ue8m0=True) + assert torch.equal(output.view(b, h * d), ref_fp8), f'{b=}, {h=}, {r=}, {d=}' + assert torch.equal(sfd.contiguous(), ref_sfd), f'{b=}, {h=}, {r=}, {d=}' + else: + output = z + deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, output) + diff = calc_diff(output, ref_z) + assert diff < (2e-3 if output_dtype == torch.float8_e4m3fn else 1e-3), \ + f'{b=}, {h=}, {r=}, {d=}, {output_dtype=}, {diff=:.5f}' 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) + t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhr,hdr->bhd', x, y, z), '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 | ' @@ -111,6 +166,7 @@ def test_fp8_bhr_hdr_bhd(use_ue8m0: bool = True): @test_filter(lambda: get_arch_major() >= 10) def test_fp8_bhd_hdr_bhr(use_ue8m0: bool = True): + deep_gemm.use_deterministic_algorithms(False) print('Testing FP8 "bhd, hdr -> bhr":') for h, r, d in [(8, 4096, 1024)]: for b in (4, 32, 128, 4096, 8192): @@ -130,7 +186,7 @@ def test_fp8_bhd_hdr_bhr(use_ue8m0: bool = True): 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) + t_cublaslt = bench_kineto(lambda: deep_gemm.einsum('bhd,hdr->bhr', x, y, z), '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 | ' @@ -154,10 +210,10 @@ def test_fp8_bhd_bhr_hdr(use_ue8m0: bool = True): 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)) + deep_gemm.fp8_einsum('bhd,bhr->hdr', x_fp8, y_fp8, z, c=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) + t = bench_kineto(lambda: deep_gemm.fp8_einsum('bhd,bhr->hdr', x_fp8, y_fp8, z, c=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 | ' @@ -175,6 +231,7 @@ def test_fp8_bhd_bhr_hdr(use_ue8m0: bool = True): test_bmk_bnk_mn() test_bhr_hdr_bhd() test_bhd_hdr_bhr() + test_bhd_bhr_hdr() test_fp8_bhr_hdr_bhd() test_fp8_bhd_hdr_bhr() diff --git a/tests/test_fp8_fp4.py b/tests/test_fp8_fp4.py index 048d70c669..eeec7911b5 100644 --- a/tests/test_fp8_fp4.py +++ b/tests/test_fp8_fp4.py @@ -1,31 +1,32 @@ -import numpy as np import random import torch import deep_gemm +from deep_gemm.utils import align, pack_ue8m0_to_int from deep_gemm.testing import ( + assert_bitwise_equal, bench_kineto, calc_diff, count_bytes, get_arch_major ) from utils import ( - assert_psum_zero_padding, + assert_direct_output_matches_fp32_accumulation, + assert_psum_zero_padding, convert_to_fp8, make_cublas_gemm, ) from generators import ( - KernelType, get_ue8m0_usage, layout_masked_to_psum, align, + KernelType, QuantConfig, get_ue8m0_usage, enumerate_normal, enumerate_m_grouped_contiguous, enumerate_m_grouped_masked, enumerate_k_grouped_contiguous, enumerate_k_grouped_contiguous_test_variants, generate_normal, generate_m_grouped_contiguous, generate_m_grouped_masked, generate_k_grouped_contiguous, - generate_k_grouped_contiguous_psum, - 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): + use_alpha_options = (False, True) if get_arch_major() == 10 else (False,) + for kernel_type, quant_config, m, n, k, major_a, major_b, accumulate, out_dtype, scores in \ + enumerate_normal(torch.float8_e4m3fn, collect_cublas_scores=True): 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' @@ -36,29 +37,53 @@ def test_gemm() -> None: 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() - 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}') + for use_alpha in use_alpha_options: + alpha = random.uniform(-1.0, 1.0) if use_alpha else None + 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, alpha=alpha) + 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() + 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, alpha=alpha) + diff = calc_diff(d, ref_d) + assert diff < quant_config.max_diff(), (f'{m=}, {n=}, {k=}, {kernel_opt}, {major_opt=}, ' + f'{accumulate=}, {out_dtype=}, {use_alpha=}, {alpha=}, ' + 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) + initial_d = d.clone() + def test_func(a_=a, b_=b, d_=d): + deep_gemm.fp8_fp4_gemm_nt(a_, b_, d_, c=d_ if accumulate else None, + disable_ue8m0_cast=disable_ue8m0_cast, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + equivalent_d = d.clone() + test_func(d_=equivalent_d) + # Bitwise deterministic test + for _ in range(20): + d.copy_(initial_d) + test_func() + assert torch.equal(d, equivalent_d), f'{m=}, {n=}, {k=}, {accumulate=}' + if quant_config.is_fp4_a or quant_config.is_fp4_b: + equivalent_fp8_d = initial_d.clone() + test_func(convert_to_fp8(a), convert_to_fp8(b), equivalent_fp8_d) + # FP4 and converted FP8 have different UMMA_K, but BF16 outputs are usually bitwise identical. + assert calc_diff(equivalent_d, equivalent_fp8_d) < 1e-14, (f'FP4/FP8 mismatch: {m=}, {n=}, {k=}, ' + f'{accumulate=}') + t = bench_kineto(test_func, 'gemm_', suppress_kineto_output=True) + cublas_func = make_cublas_gemm(a, b, d, c) + cublas_times = bench_kineto(cublas_func, ('nvjet', 'bstensorop', 'reduce'), + suppress_kineto_output=True, with_multiple_kernels=True) + cublas_t = sum(cublas_times) + if cublas_t > 0: + scores.append(cublas_t / t) 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") + f'{cublas_t / t:.2f}x cuBLAS speedup') def test_m_grouped_gemm_contiguous() -> None: @@ -77,30 +102,36 @@ def test_m_grouped_gemm_contiguous() -> None: 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) + m, a, b, grouped_layout, d, ref_d, valid_mask = 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, ensure_zero_padding=ensure_zero_padding, - 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}, ' - f'{diff:.5f}, alias={test_alias}, {ensure_zero_padding=}') - if ensure_zero_padding: - assert_psum_zero_padding(a, d, grouped_layout, 'FP8/FP4') - 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, + def test_func(a_=a, b_=b, d_=d): + getattr(deep_gemm, func_name)(a_, b_, d_, grouped_layout, disable_ue8m0_cast=disable_ue8m0_cast, + use_psum_layout=use_psum_layout, ensure_zero_padding=ensure_zero_padding, + recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) + equivalent_d = d.clone() + test_func(d_=equivalent_d) + # Bitwise deterministic test + for _ in range(20): + test_func() + assert torch.equal(d[valid_mask], equivalent_d[valid_mask]), f'{m=}, {n=}, {k=}, alias={test_alias}' + if quant_config.is_fp4_a or quant_config.is_fp4_b: + equivalent_fp8_d = d.clone() + test_func(convert_to_fp8(a), convert_to_fp8(b), equivalent_fp8_d) + # FP4 and converted FP8 have different UMMA_K, but BF16 outputs are usually bitwise identical. + assert calc_diff(equivalent_d[valid_mask], equivalent_fp8_d[valid_mask]) < 1e-14, ( + f'FP4/FP8 mismatch: {m=}, {n=}, {k=}, alias={test_alias}') + diff = calc_diff(d[valid_mask], ref_d[valid_mask]) + assert diff < quant_config.max_diff(), (f'{m=}, {n=}, {k=}, {major_opt}, {kernel_opt}, ' + f'{diff:.5f}, alias={test_alias}, {ensure_zero_padding=}') + if use_psum_layout and ensure_zero_padding: + assert_psum_zero_padding(a, d, grouped_layout, 'FP8/FP4') + m, a, b, grouped_layout, d, ref_d, valid_mask = 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) @@ -132,42 +163,41 @@ def test_m_grouped_gemm_masked() -> None: num_tests = 8 sum_t, max_t = 0, 0 sum_ops, sum_bytes = 0, 0 + expected_m = int(expected_m_per_group * 1.2) # Select best alignment - alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(int(expected_m_per_group * 1.2)) + alignment = deep_gemm.get_theoretical_mk_alignment_for_contiguous_layout(expected_m) 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 + a, b, grouped_layout, d, ref_d, valid_mask = 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) + def test_func(a_=a, b_=b, d_=d): + common = dict(disable_ue8m0_cast=disable_ue8m0_cast, recipe=recipe, recipe_a=recipe_a, recipe_b=recipe_b) 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]] + deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(a_, b_, d_, grouped_layout, + use_psum_layout=True, expected_m_for_psum_layout=expected_m, **common) 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}' + deep_gemm.m_grouped_fp8_fp4_gemm_nt_masked(a_, b_, d_, grouped_layout, expected_m, **common) + + equivalent_d = d.clone() + test_func(d_=equivalent_d) + # Bitwise deterministic test + for _ in range(20): + test_func() + assert torch.equal(d[valid_mask], equivalent_d[valid_mask]), f'{max_m=}, {n=}, {k=}' + if quant_config.is_fp4_a or quant_config.is_fp4_b: + equivalent_fp8_d = d.clone() + test_func(convert_to_fp8(a), convert_to_fp8(b), equivalent_fp8_d) + # FP4 and converted FP8 have different UMMA_K, but BF16 outputs are usually bitwise identical. + assert calc_diff(equivalent_d[valid_mask], equivalent_fp8_d[valid_mask]) < 1e-14, ( + f'FP4/FP8 mismatch: {max_m=}, {n=}, {k=}, {num_groups=}') + diff = calc_diff(d[valid_mask], ref_d[valid_mask]) + assert diff < quant_config.max_diff(), f'{max_m=}, {n=}, {k=}, {kernel_opt}, {num_groups=}, {diff:.5f}' # Test performance with fixed shapes - valid_m = masked_m.sum().item() + valid_m = int(valid_mask.sum().item()) t = bench_kineto(test_func, 'gemm_', suppress_kineto_output=True) sum_t += t @@ -184,54 +214,142 @@ def test_func(): 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, real_ks_cpu, aligned_ks_cpu, _, gran_k, k_alignment, use_psum_layout in enumerate_k_grouped_contiguous(torch.float8_e4m3fn): - recipe = (1, 1, gran_k) - use_ue8m0 = get_ue8m0_usage(KernelType.Kernel1D1D) - - for test_real_ks_cpu, test_aligned_ks_cpu, _, _ in enumerate_k_grouped_contiguous_test_variants(real_ks_cpu, k_alignment, use_psum_layout): - if use_psum_layout: - total_k, a, b, c, d, ref_d, grouped_layout, _ = generate_k_grouped_contiguous_psum(num_groups, m, n, major_a, major_b, test_real_ks_cpu, k_alignment=k_alignment, use_ue8m0=use_ue8m0, gran_k=gran_k) - else: - total_k, a, b, c, d, ref_d, grouped_layout, _ = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, test_aligned_ks_cpu, use_ue8m0=use_ue8m0, gran_k=gran_k) - c_orig = c.clone() if use_psum_layout else None - k_grouped_fp8_gemm_contiguous(a, b, d, test_aligned_ks_cpu, grouped_layout, c, recipe=recipe, use_psum_layout=use_psum_layout) - - diff = calc_diff(d, ref_d) - assert diff < 0.001, f'{m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, {test_aligned_ks_cpu=}, {use_psum_layout=}, {diff:.5f}' - - # Unsynced psum paths - if use_psum_layout: - c.copy_(c_orig) - k_grouped_fp8_gemm_contiguous(a, b, d, None, grouped_layout, c, recipe=recipe, - use_psum_layout=True) - diff = calc_diff(d, ref_d) - assert diff < 0.001, f'None ks_cpu path: {m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, {diff:.5f}' + print('Testing k-grouped GEMM:') + + arch_major = get_arch_major() + test_options = [(torch.float8_e4m3fn, QuantConfig(), + deep_gemm.k_grouped_fp8_gemm_nt_contiguous if arch_major == 9 else + deep_gemm.k_grouped_fp8_gemm_tn_contiguous)] + if arch_major == 10: + test_options.append((torch.float4_e2m1fn_x2, QuantConfig((32, 32, True, True)), + deep_gemm.k_grouped_fp4_gemm_nt_contiguous)) + use_ue8m0 = get_ue8m0_usage(KernelType.Kernel1D1D) + for dtype, quant_config, gemm in test_options: + is_fp4 = dtype == torch.float4_e2m1fn_x2 + dtype_opt = 'FP4' if is_fp4 else 'FP8' + + for num_groups, m, n, major_a, major_b, real_ks_cpu, _, _, gran_k, k_alignment, use_psum_layout, accumulate, out_dtype in \ + enumerate_k_grouped_contiguous(dtype): + recipe = (1, 1, gran_k) + + for test_real_ks_cpu in enumerate_k_grouped_contiguous_test_variants(real_ks_cpu): + total_k, a, b, c, d, ref_d, grouped_layout, host_ks_cpu = generate_k_grouped_contiguous( + num_groups, m, n, major_a, major_b, test_real_ks_cpu, + use_ue8m0=use_ue8m0, gran_k=gran_k, + quant_config=quant_config if is_fp4 else None, + use_psum_layout=use_psum_layout, k_alignment=k_alignment, + accumulate=accumulate, out_dtype=out_dtype) + + initial_d = d.clone() + if not accumulate: + initial_d.fill_(float('nan')) + equivalent_d = initial_d.clone() + gemm(a, b, equivalent_d, host_ks_cpu, grouped_layout, equivalent_d if accumulate else None, + recipe=recipe, use_psum_layout=use_psum_layout) + if is_fp4 and accumulate: + fp8_a, fp8_b = convert_to_fp8(a), convert_to_fp8(b) + fp8_a = (fp8_a[0].T.contiguous(), fp8_a[1]) + fp8_b = (fp8_b[0].T.contiguous(), fp8_b[1]) + equivalent_fp8_d = initial_d.clone() + deep_gemm.k_grouped_fp8_gemm_tn_contiguous( + fp8_a, fp8_b, equivalent_fp8_d, host_ks_cpu, grouped_layout, equivalent_fp8_d if accumulate else None, + recipe=recipe, use_psum_layout=use_psum_layout) + mismatch_message = (f'FP4/FP8 mismatch: {m=}, {n=}, {total_k=}, ' + f'{test_real_ks_cpu=}, {use_psum_layout=}') + assert calc_diff(equivalent_d, equivalent_fp8_d) < 1e-14, mismatch_message + + # Bitwise deterministic test + host_ks_options = (host_ks_cpu, None, []) if use_psum_layout else (host_ks_cpu, ) + for test_host_ks_cpu in host_ks_options: + for stress_idx in range(20): + d.copy_(initial_d) + gemm(a, b, d, test_host_ks_cpu, grouped_layout, c, + recipe=recipe, use_psum_layout=use_psum_layout) + assert_bitwise_equal( + d, equivalent_d, + f'k-grouped self-consistency at {stress_idx=}, {dtype_opt}, {m=}, {n=}, {total_k=}, ' + f'{test_real_ks_cpu=}, {test_host_ks_cpu=}, {use_psum_layout=}, {accumulate=}, {out_dtype=}' + ) + if not accumulate: + case_label = (f'{dtype_opt} K-grouped direct output, {m=}, {n=}, {total_k=}, ' + f'{test_real_ks_cpu=}, {test_host_ks_cpu=}, {use_psum_layout=}, ' + f'{out_dtype=}') + assert_direct_output_matches_fp32_accumulation( + d, + lambda output, accumulator: gemm( + a, b, output, test_host_ks_cpu, grouped_layout, accumulator, + recipe=recipe, use_psum_layout=use_psum_layout), + case_label) + + if accumulate: + diff = calc_diff(d, ref_d) + assert diff < quant_config.max_diff(), ( + f'{dtype_opt}, {m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, ' + f'{host_ks_cpu=}, {use_psum_layout=}, {accumulate=}, {out_dtype=}, {diff:.5f}') + + # gran_k=128 requires FP32 SF input; only gran_k=32 accepts + # the per-group packed INT32 UE8M0 layout. + if gran_k == 32: + sf_ks = [k // gran_k for k in host_ks_cpu] + ref_packed_a = torch.cat([ + pack_ue8m0_to_int(torch.nn.functional.pad( + group_sf.T, (0, align(sf_k, 4) - sf_k)).contiguous()).T + for group_sf, sf_k in zip(a[1].split(sf_ks), sf_ks) if sf_k > 0 + ]) + ref_packed_b = torch.cat([ + pack_ue8m0_to_int(torch.nn.functional.pad( + group_sf.T, (0, align(sf_k, 4) - sf_k)).contiguous()).T + for group_sf, sf_k in zip(b[1].split(sf_ks), sf_ks) if sf_k > 0 + ]) + packed_a = ( + a[0], deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + a[1], grouped_layout, host_ks_cpu, gran_k, k_alignment, use_psum_layout)) + packed_b = ( + b[0], deep_gemm.get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + b[1], grouped_layout, host_ks_cpu, gran_k, k_alignment, use_psum_layout)) + assert torch.equal(packed_a[1], ref_packed_a) + assert torch.equal(packed_b[1], ref_packed_b) + + packed_d = initial_d.clone() + gemm(packed_a, packed_b, packed_d, host_ks_cpu, grouped_layout, packed_d if accumulate else None, + recipe=recipe, use_psum_layout=use_psum_layout) + if accumulate: + packed_diff = calc_diff(packed_d, ref_d) + assert packed_diff < quant_config.max_diff(), ( + f'pre-packed INT32 SF: {dtype_opt}, {m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, ' + f'{host_ks_cpu=}, {use_psum_layout=}, {accumulate=}, {out_dtype=}, ' + f'{packed_diff:.5f}') + else: + case_label = (f'pre-packed INT32 SF direct output: {dtype_opt}, {m=}, {n=}, ' + f'{total_k=}, {test_real_ks_cpu=}, {host_ks_cpu=}, ' + f'{use_psum_layout=}, {out_dtype=}') + assert_direct_output_matches_fp32_accumulation( + packed_d, + lambda output, accumulator: gemm( + packed_a, packed_b, output, host_ks_cpu, grouped_layout, accumulator, + recipe=recipe, use_psum_layout=use_psum_layout), + case_label) + + _, a, b, c, d, _, grouped_layout, host_ks_cpu = generate_k_grouped_contiguous( + num_groups, m, n, major_a, major_b, real_ks_cpu, + use_ue8m0=use_ue8m0, gran_k=gran_k, + quant_config=quant_config if is_fp4 else None, + use_psum_layout=use_psum_layout, k_alignment=k_alignment, + accumulate=accumulate, out_dtype=out_dtype) - c.copy_(c_orig) - k_grouped_fp8_gemm_contiguous(a, b, d, [], grouped_layout, c, recipe=recipe, - use_psum_layout=True) - diff = calc_diff(d, ref_d) - assert diff < 0.001, f'empty ks_cpu path: {m=}, {n=}, {total_k=}, {test_real_ks_cpu=}, {diff:.5f}' - - # Test performance - if use_psum_layout: - total_k, a, b, c, d, ref_d, grouped_layout, _ = generate_k_grouped_contiguous_psum(num_groups, m, n, major_a, major_b, real_ks_cpu, k_alignment=k_alignment, use_ue8m0=use_ue8m0, gran_k=gran_k) - else: - total_k, a, b, c, d, ref_d, grouped_layout, _ = generate_k_grouped_contiguous(num_groups, m, n, major_a, major_b, aligned_ks_cpu, use_ue8m0=use_ue8m0, gran_k=gran_k) - - # noinspection PyShadowingNames - def test_func(): - k_grouped_fp8_gemm_contiguous(a, b, d, aligned_ks_cpu, grouped_layout, c, recipe=recipe, use_psum_layout=use_psum_layout) + # noinspection PyShadowingNames + def test_func(): + gemm(a, b, d, host_ks_cpu, grouped_layout, c, + recipe=recipe, use_psum_layout=use_psum_layout) - t = bench_kineto(test_func, 'gemm_', suppress_kineto_output=True) - print(f' > Perf ({num_groups=:2}, m={m:5}, n={n:5}, k={total_k:5}, gran_k={gran_k:3}, k_alignment={k_alignment:3}, psum={int(use_psum_layout)}): ' - f'{t * 1e6:4.0f} us | ' - f'{2 * m * n * total_k / t / 1e12:4.0f} TFLOPS | ' - f'{count_bytes(a, b, c, d) / 1e9 / t:4.0f} GB/s') + t = bench_kineto(test_func, 'gemm_', suppress_kineto_output=True) + logical_k = sum(real_ks_cpu) + out_opt = 'FP32' if out_dtype == torch.float else 'BF16' + print(f' > Perf ({dtype_opt}, {num_groups=:2}, m={m:5}, n={n:5}, k={logical_k:5}, gran_k={gran_k:3}, ' + f'k_alignment={k_alignment:3}, psum={int(use_psum_layout)}, acc={int(accumulate)}, {out_opt}): ' + f'{t * 1e6:4.0f} us | ' + f'{2 * m * n * logical_k / t / 1e12:4.0f} TFLOPS | ' + f'{count_bytes(a, b, c, d) / 1e9 / t:4.0f} GB/s') print() diff --git a/tests/test_layout.py b/tests/test_layout.py index fe6c5d3247..8fdcaa8ae7 100644 --- a/tests/test_layout.py +++ b/tests/test_layout.py @@ -82,24 +82,34 @@ def test_sf_layout_kernels() -> None: def test_k_grouped_sf_layout_kernels() -> None: print('Testing k-grouped SF layout kernels:') - for mn, ks_cpu, num_groups, gran_k in enumerate_k_grouped_sf_layout(): + if get_arch_major() != 10: + print(' > Skipped (packed k-grouped SF only supported on SM100)') + return + + for mn, ks_cpu, num_groups, gran_k, k_alignment in enumerate_k_grouped_sf_layout(): sf_ks = [k // gran_k for k in ks_cpu] - packed_sf_ks = [ceil_div(k, gran_k * 4) for k in ks_cpu] + packed_sf_ks = [ceil_div(sf_k, 4) for sf_k in sf_ks] grouped_layout = torch.tensor(ks_cpu, dtype=torch.int, device='cuda') x = torch.randn((sum(ks_cpu), mn), dtype=torch.bfloat16, device='cuda') - x, fp32_sf = per_channel_cast_to_fp8(x, use_ue8m0=True, gran_k=gran_k) + _, 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, grouped_layout, ks_cpu, gran_k, gran_k) + # Correctness: each group starts at a new packed row and independently + # zero-fills unused slots in its final row. + packed_sf = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, grouped_layout, ks_cpu, gran_k, k_alignment) + assert packed_sf.shape == (sum(packed_sf_ks), mn) 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=}' + for group_idx in range(num_groups): + ref_packed_sf = get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(split_fp32_sf[group_idx].T).T + assert torch.equal(split_packed_sf[group_idx], ref_packed_sf), f'{group_idx=}' + if sf_ks[group_idx] % 4: + packed_bytes = split_packed_sf[group_idx][-1].view(torch.uint8).view(mn, 4) + assert torch.count_nonzero(packed_bytes[:, sf_ks[group_idx] % 4:]) == 0 # Performance - t = bench_kineto(lambda: get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, grouped_layout, ks_cpu, gran_k, gran_k), 'pack_fp32_into_ue8m0') - print(f' > Perf ({num_groups=:3}, {mn=:5}, sum_k={sum(ks_cpu):5}, gran_k={gran_k:3}):' + t = bench_kineto(lambda: get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, grouped_layout, ks_cpu, gran_k, k_alignment), 'pack_fp32_into_ue8m0') + print(f' > Perf ({num_groups=:3}, {mn=:5}, sum_k={sum(ks_cpu):5}, ' + f'gran_k={gran_k:3}, k_alignment={k_alignment:3}): ' f'{t * 1e6:4.0f} us | ' f'{count_bytes(fp32_sf, packed_sf, grouped_layout) / 1e9 / t:4.0f} GB/s') print() @@ -108,39 +118,33 @@ def test_k_grouped_sf_layout_kernels() -> None: def test_k_grouped_psum_sf_layout_kernels() -> None: print('Testing k-grouped psum SF layout kernels:') if get_arch_major() != 10: - print(' > Skipped (psum SF pack kernel only supported on SM100)') + print(' > Skipped (packed k-grouped SF only supported on SM100)') return for mn, real_ks_cpu, aligned_ks_cpu, psum_layout, num_groups, gran_k, k_alignment in enumerate_k_grouped_psum_sf_layout(): grouped_layout = torch.tensor(psum_layout, dtype=torch.int, device='cuda') - # The pack kernel only consumes SF values, so generate random UE8M0 SF directly - # (compact layout: `ceil_div(k, gran_k)` rows per group) - num_sf_rows = sum(ceil_div(k, gran_k) for k in real_ks_cpu) + sf_ks = [align(k, k_alignment) // gran_k for k in real_ks_cpu] + packed_sf_ks = [ceil_div(sf_k, 4) for sf_k in sf_ks] + num_sf_rows = sum(sf_ks) fp32_sf = ceil_to_ue8m0(torch.rand((num_sf_rows, mn), dtype=torch.float, device='cuda') + 0.5) + ref_packed_sf = torch.cat([ + get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(group_sf.T).T + for group_sf in fp32_sf.split(sf_ks) if group_sf.size(0) > 0 + ]) - ref_packed_sf = [] - sf_start = 0 - for k in real_ks_cpu: - sf_end = sf_start + ceil_div(k, gran_k) - ref_packed_sf.append(get_mn_major_tma_aligned_packed_ue8m0_tensor_torch_impl(fp32_sf[sf_start:sf_end].T).T) - sf_start = sf_end - ref_packed_sf = torch.cat(ref_packed_sf) - - exact_packed_sf = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, grouped_layout, real_ks_cpu, gran_k, k_alignment, use_psum_layout=True) - assert torch.equal(exact_packed_sf, ref_packed_sf) - - # Aligned K sizes match the GEMM API path and may allocate upper-bound rows. + # Psum layout provides raw K; SF covers each group's aligned K span and + # packing independently pads every group to a packed row boundary. packed_sf = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, grouped_layout, aligned_ks_cpu, gran_k, k_alignment, use_psum_layout=True) - assert torch.equal(packed_sf[:ref_packed_sf.size(0)], ref_packed_sf) + assert torch.equal(packed_sf, ref_packed_sf) - # Unsynced upper-bound paths - upper_bound_packed_sf = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( + # Without synchronized aligned K sizes, allocation uses a safe upper bound. + unsynced_packed_sf = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( fp32_sf, grouped_layout, None, gran_k, k_alignment, use_psum_layout=True) - assert torch.equal(upper_bound_packed_sf[:ref_packed_sf.size(0)], ref_packed_sf) + assert torch.equal(unsynced_packed_sf[:sum(packed_sf_ks)], ref_packed_sf) empty_ks_packed_sf = get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor( fp32_sf, grouped_layout, [], gran_k, k_alignment, use_psum_layout=True) - assert torch.equal(empty_ks_packed_sf[:ref_packed_sf.size(0)], ref_packed_sf) + assert torch.equal(empty_ks_packed_sf[:sum(packed_sf_ks)], ref_packed_sf) # Performance t = bench_kineto(lambda: get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(fp32_sf, grouped_layout, aligned_ks_cpu, gran_k, k_alignment, use_psum_layout=True), 'pack_fp32_into_ue8m0') diff --git a/tests/test_legacy.py b/tests/test_legacy.py index 587798075a..525f5579e2 100644 --- a/tests/test_legacy.py +++ b/tests/test_legacy.py @@ -19,7 +19,7 @@ def test_m_grouped_gemm_contiguous_tl() -> None: 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) + 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() @@ -33,7 +33,7 @@ def test_m_grouped_gemm_contiguous_tl() -> None: 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) + 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(): @@ -49,9 +49,10 @@ def test_func(): def test_k_grouped_gemm_contiguous_tl() -> None: print('Testing k-grouped contiguous Triton GEMM:') - for num_groups, m, n, major_a, major_b, _, aligned_ks_cpu, expected_k_per_group, _, k_alignment, use_psum_layout in enumerate_k_grouped_contiguous(torch.bfloat16): + for num_groups, m, n, major_a, major_b, _, aligned_ks_cpu, expected_k_per_group, _, k_alignment, use_psum_layout, accumulate, out_dtype in \ + enumerate_k_grouped_contiguous(torch.bfloat16): # Legacy Triton kernels do not mask per-group K tails inside the BLOCK_SIZE_K loop. - if use_psum_layout or k_alignment != 128: + if use_psum_layout or k_alignment != 128 or not accumulate or out_dtype != torch.float: continue major_opt = 'N' if major_a.is_k_major() else 'T' diff --git a/tests/test_mega_gate.py b/tests/test_mega_gate.py new file mode 100644 index 0000000000..3d282f3077 --- /dev/null +++ b/tests/test_mega_gate.py @@ -0,0 +1,354 @@ +from typing import Any + +import torch +import torch.nn.functional as F + +import deep_gemm +from deep_gemm.testing import ( + assert_bitwise_equal, + bench_kineto, + calc_diff, + count_bytes, + get_arch_major, + test_filter, +) +from tile_kernels.moe import moe_topk_gate + +BASELINE_KERNELS = ('nvjet', 'splitKreduce', 'cutlass', 'moe_topk_gate') + +NUM_TOPK = 6 +NUM_SHARED_EXPERTS = 1 +ROUTED_SCALING_FACTOR = 1.5 +VLLM_NUM_EXTRA_EXPERTS = 32 + +GATE_CONFIGS = ( + (64, 2, 6), (72, 1, 6), (96, 2, 6), (108, 2, 6), (128, 2, 6), (144, 2, 6), + (256, 1, 6), (256, 2, 8), (384, 1, 6), (384, 1, 9), (512, 1, 6), +) + + +def enumerate_perf_cases(): + yield from ((num_tokens, hidden, 384) + for num_tokens in (1, 3, 16, 128, 512, 1024, 2048, 4096, 8192, 16384, 32768) + for hidden in (4096, 5120, 7168)) + yield from ( + (512, 4096, 128), (1024, 4096, 128), (4096, 5120, 128), (32768, 7168, 128), + (128, 4096, 256), (8192, 4352, 264), (4200, 7168, 256), (8192, 7168, 256), + (128, 7168, 512), (8192, 4096, 512), (32768, 7168, 512), + ) + + +def make_physical_map(num_logical_experts: int) -> tuple[torch.Tensor, torch.Tensor]: + to_physical_map = torch.full( + (num_logical_experts, VLLM_NUM_EXTRA_EXPERTS + 1), -1, dtype=torch.int32, device='cuda') + logical_count = torch.ones((num_logical_experts,), dtype=torch.int32, device='cuda') + to_physical_map[:, 0] = torch.arange(num_logical_experts, dtype=torch.int32, device='cuda') + num_duplicated_experts = min(num_logical_experts, VLLM_NUM_EXTRA_EXPERTS) + to_physical_map[:num_duplicated_experts, 1] = ( + num_logical_experts + torch.arange(num_duplicated_experts, dtype=torch.int32, device='cuda')) + logical_count[:num_duplicated_experts] = 2 + return to_physical_map, logical_count + + +def make_inputs(num_tokens: int, hidden: int, num_routed_experts: int, + use_shared_as_routed: bool = False, with_bias: bool = True, + with_image_bias: bool = False, with_mask: bool = False, + with_physical_map: bool = True, + num_shared_experts: int = NUM_SHARED_EXPERTS) -> dict[str, Any]: + num_logical_experts = num_routed_experts + (num_shared_experts if use_shared_as_routed else 0) + mask = torch.rand((num_tokens,), device='cuda') >= 0.25 if with_mask else None + if mask is not None: + mask[0] = True + to_physical_map, logical_count = make_physical_map(num_logical_experts) if with_physical_map else (None, None) + return dict( + x=torch.randn((num_tokens, hidden), dtype=torch.bfloat16, device='cuda'), + weight=torch.randn((num_routed_experts, hidden), dtype=torch.bfloat16, device='cuda').mul_(hidden ** -0.5), + mask=mask, + bias=torch.randn((num_routed_experts,), dtype=torch.float32, device='cuda') if with_bias else None, + image_bias=torch.randn((num_routed_experts,), dtype=torch.float32, device='cuda') if with_image_bias else None, + image_token_mask=torch.rand((num_tokens,), device='cuda') < 0.5 if with_image_bias else None, + to_physical_map=to_physical_map, + logical_count=logical_count, + ) + + +def make_outputs(inputs: dict[str, Any], use_shared_as_routed: bool = False, + with_unmapped_topk_idx: bool = False, num_topk: int = NUM_TOPK, + num_shared_experts: int = NUM_SHARED_EXPERTS) -> dict[str, torch.Tensor | None]: + shape = (inputs['x'].shape[0], num_topk + (num_shared_experts if use_shared_as_routed else 0)) + return dict( + topk_idx=torch.empty(shape, dtype=torch.int64, device='cuda'), + topk_weights=torch.empty(shape, dtype=torch.float32, device='cuda'), + unmapped_topk_idx=(torch.empty((shape[0], num_topk), dtype=torch.int64, device='cuda') + if with_unmapped_topk_idx else None), + ) + + +def run_mega_gate(inputs: dict[str, Any], outputs: dict[str, torch.Tensor | None], scoring_func: str, + use_shared_as_routed: bool = False, + num_topk: int = NUM_TOPK, num_shared_experts: int = NUM_SHARED_EXPERTS, + **overrides: Any) -> tuple[torch.Tensor, torch.Tensor]: + kwargs = {**inputs, **overrides} + x, weight = kwargs.pop('x'), kwargs.pop('weight') + unmapped_topk_idx = kwargs.pop('unmapped_topk_idx', outputs['unmapped_topk_idx']) + return deep_gemm.bf16_mega_gate( + x, weight, num_topk, use_shared_as_routed, num_shared_experts, ROUTED_SCALING_FACTOR, 0, + scoring_func=scoring_func, unmapped_topk_idx=unmapped_topk_idx, + out=(outputs['topk_idx'], outputs['topk_weights']), **kwargs) + + +def run_reference_gemm(inputs: dict[str, Any]) -> torch.Tensor: + logits = torch.empty( + (inputs['x'].shape[0], inputs['weight'].shape[0]), dtype=torch.float32, device='cuda') + deep_gemm.cublaslt_gemm_nt(inputs['x'], inputs['weight'], logits) + return logits + + +def apply_scoring(logits: torch.Tensor, scoring_func: str) -> torch.Tensor: + assert scoring_func in ('sigmoid', 'sqrtsoftplus', 'identity') + return (torch.sigmoid(logits) if scoring_func == 'sigmoid' else + F.softplus(logits).sqrt() if scoring_func == 'sqrtsoftplus' else logits) + + +def check_topk_weights(inputs: dict[str, Any], outputs: dict[str, torch.Tensor | None], logits: torch.Tensor, + scoring_func: str, + mask: torch.Tensor | None = None, + force_random: torch.Tensor | None = None, + fix_routing_mask: torch.Tensor | None = None, + num_topk: int = NUM_TOPK) -> float: + fused_idx, fused_weights = outputs['topk_idx'], outputs['topk_weights'] + num_tokens, num_routed_experts = logits.shape + active = mask if mask is not None else torch.ones((num_tokens,), dtype=torch.bool, device='cuda') + score_based = active if force_random is None else active & ~force_random + weight_diff = 0.0 + + if score_based.any(): + scores = apply_scoring(logits, scoring_func) + logical_topk_idx = outputs['unmapped_topk_idx'] + if logical_topk_idx is None: + logical_topk_idx = fused_idx[:, :num_topk] + to_physical_map = inputs['to_physical_map'] + if to_physical_map is not None: + logical_topk_idx = torch.where( + logical_topk_idx < to_physical_map.shape[0], + logical_topk_idx, logical_topk_idx - to_physical_map.shape[0]) + + selected_idx = logical_topk_idx[score_based] + assert torch.all((selected_idx >= 0) & (selected_idx < num_routed_experts)).item() + + ranking_bias = (inputs['bias'] if inputs['bias'] is not None + else torch.zeros((num_routed_experts,), dtype=torch.float32, device='cuda')) + if inputs['image_token_mask'] is not None: + ranking_bias = torch.where( + inputs['image_token_mask'][:, None], inputs['image_bias'], ranking_bias) + ranking = scores + ranking_bias + + def reference_weights(row_mask: torch.Tensor, topk_idx: torch.Tensor) -> torch.Tensor: + topk_scores = scores[row_mask].gather(1, topk_idx) + return topk_scores / (topk_scores.sum(dim=1, keepdim=True) + 1e-20) * ROUTED_SCALING_FACTOR + + weight_diff = float(calc_diff(fused_weights[score_based, :num_topk], + reference_weights(score_based, selected_idx))) + assert weight_diff < 1e-11, f'{weight_diff=}' + + ranked = score_based if fix_routing_mask is None else score_based & ~fix_routing_mask + if ranked.any(): + ranked_idx = logical_topk_idx[ranked] + sorted_idx = ranked_idx.sort(dim=1).values + assert torch.all(sorted_idx[:, 1:] != sorted_idx[:, :-1]).item() + + selected_ranking = ranking[ranked].gather(1, ranked_idx) + unselected_max = ranking[ranked].scatter(1, ranked_idx, float('-inf')).max(dim=1).values + crossing = unselected_max - selected_ranking.min(dim=1).values + tolerance = 1e-5 + 1e-5 * ranking[ranked].abs().amax(dim=1) + assert torch.all(crossing <= tolerance).item(), \ + f'top-k selection beyond numerical ties: {crossing.max().item()=}' + + ref_weights = reference_weights(ranked, ranking[ranked].topk(num_topk, dim=1).indices) + ranked_diff = float(calc_diff(fused_weights[ranked, :num_topk].sort(dim=1).values, + ref_weights.sort(dim=1).values)) + assert ranked_diff < 1e-11, f'{ranked_diff=}' + weight_diff = max(weight_diff, ranked_diff) + + if fused_weights.shape[1] > num_topk: + assert torch.all(fused_weights[score_based, num_topk:] == 1.0).item() + + if (~active).any(): + assert torch.all(fused_idx[~active] == -1).item() + assert torch.all(fused_weights[~active] == 0).item() + if outputs['unmapped_topk_idx'] is not None: + assert torch.all(outputs['unmapped_topk_idx'][~active] == -1).item() + + return weight_diff + + +def check_correctness(inputs: dict[str, Any], outputs: dict[str, torch.Tensor | None], scoring_func: str, + use_shared_as_routed: bool = False, + num_topk: int = NUM_TOPK, + num_shared_experts: int = NUM_SHARED_EXPERTS) -> float: + logits = run_reference_gemm(inputs) + run_mega_gate( + inputs, outputs, scoring_func, use_shared_as_routed, + num_topk=num_topk, num_shared_experts=num_shared_experts) + return check_topk_weights( + inputs, outputs, logits, scoring_func, + inputs['mask'], num_topk=num_topk) + + +def check_case(num_tokens: int, hidden: int, num_routed_experts: int, + scoring_func: str = 'sqrtsoftplus', use_shared_as_routed: bool = False, + num_topk: int = NUM_TOPK, num_shared_experts: int = NUM_SHARED_EXPERTS, + with_unmapped_topk_idx: bool = False, **input_options: Any) -> None: + inputs = make_inputs(num_tokens, hidden, num_routed_experts, use_shared_as_routed, + num_shared_experts=num_shared_experts, **input_options) + outputs = make_outputs(inputs, use_shared_as_routed, with_unmapped_topk_idx, num_topk, num_shared_experts) + check_correctness(inputs, outputs, scoring_func, use_shared_as_routed, num_topk, num_shared_experts) + + +def bench_baseline_gate(inputs: dict[str, Any]) -> float: + outputs = make_outputs(inputs) + logits = torch.empty( + (inputs['x'].shape[0], inputs['weight'].shape[0]), dtype=torch.float32, device='cuda') + gate_kwargs = {name: inputs[name] for name in + ('mask', 'bias', 'image_bias', 'image_token_mask', 'to_physical_map', 'logical_count')} + + def run_baseline(): + deep_gemm.cublaslt_gemm_nt(inputs['x'], inputs['weight'], logits) + return moe_topk_gate( + logits, NUM_TOPK, False, NUM_SHARED_EXPERTS, ROUTED_SCALING_FACTOR, 0, + scoring_func='sqrtsoftplus', out=(outputs['topk_idx'], outputs['topk_weights']), + unmapped_topk_idx=None, **gate_kwargs) + + return sum(bench_kineto(run_baseline, BASELINE_KERNELS, suppress_kineto_output=True)) + + +@test_filter(lambda: get_arch_major() == 10) +@torch.no_grad() +def test_mega_gate() -> None: + torch.manual_seed(0) + + print('Testing BF16 GEMM + MoE top-k gate:') + for num_tokens, hidden, num_routed_experts in enumerate_perf_cases(): + inputs = make_inputs(num_tokens, hidden, num_routed_experts) + outputs = make_outputs(inputs) + weight_diff = check_correctness(inputs, outputs, 'sqrtsoftplus') + + fused_t = bench_kineto( + lambda: run_mega_gate(inputs, outputs, 'sqrtsoftplus'), + 'mega_gate', suppress_kineto_output=True) + baseline_t = bench_baseline_gate(inputs) + num_flops = 2 * num_tokens * num_routed_experts * hidden + num_bytes = count_bytes(*inputs.values(), *outputs.values()) + print(f' > Perf (m={num_tokens:6}, n={num_routed_experts:4}, k={hidden:4}, topk={NUM_TOPK}, ' + f'sqrtsoftplus): ' + f'{fused_t * 1e6:7.1f} us | ' + f'{num_flops / fused_t / 1e12:4.0f} TFLOPS | ' + f'{num_bytes / fused_t / 1e9:4.0f} GB/s | ' + f'baseline {baseline_t * 1e6:7.1f} us, {baseline_t / fused_t:5.2f}x | ' + f'diff {weight_diff:.1e}') + print() + + +def check_routing_controls() -> None: + torch.manual_seed(1) + num_tokens, hidden, num_routed_experts = 512, 4096, 384 + inputs = make_inputs(num_tokens, hidden, num_routed_experts, use_shared_as_routed=True, with_image_bias=True) + outputs = make_outputs(inputs, use_shared_as_routed=True, with_unmapped_topk_idx=True) + logits = run_reference_gemm(inputs) + + token_idx = torch.arange(num_tokens, device='cuda') + mask, fix_routing_mask, force_random = token_idx % 11 != 0, token_idx % 3 == 0, token_idx % 5 == 0 + fixed_topk_idx = torch.randint( + 0, num_routed_experts, (num_tokens, NUM_TOPK), dtype=torch.int64, device='cuda') + unmapped_topk_idx = outputs['unmapped_topk_idx'] + unmapped_topk_idx.copy_(fixed_topk_idx) + + fused_idx, fused_weights = run_mega_gate( + inputs, outputs, 'sqrtsoftplus', True, mask=mask, + fix_routing_mask=fix_routing_mask, force_random=force_random) + check_topk_weights( + inputs, outputs, logits, 'sqrtsoftplus', mask, force_random, fix_routing_mask) + + active_force_random = force_random & mask + fixed = fix_routing_mask & mask & ~force_random + assert torch.equal(unmapped_topk_idx[fixed], fixed_topk_idx[fixed]) + random_idx, random_weights = fused_idx[active_force_random], fused_weights[active_force_random] + assert torch.all((random_idx >= 0) & (random_idx < inputs['logical_count'].sum())).item() + assert torch.all(random_weights > 0).item() + sorted_random_idx = random_idx.sort(dim=1).values + assert torch.all(sorted_random_idx[:, 1:] != sorted_random_idx[:, :-1]).item() + assert torch.all(unmapped_topk_idx[active_force_random] == -1).item() + + +@test_filter(lambda: get_arch_major() == 10) +@torch.no_grad() +def test_mega_gate_api_contract() -> None: + torch.manual_seed(2) + + for num_routed_experts, num_shared_experts, num_topk in GATE_CONFIGS: + can_route_shared = ( + num_topk % num_shared_experts == 0 and + num_routed_experts % (num_topk // num_shared_experts) == 0) + for use_shared_as_routed in ((False, True) if can_route_shared else (False,)): + check_case(16, 4096, num_routed_experts, use_shared_as_routed=use_shared_as_routed, + num_topk=num_topk, num_shared_experts=num_shared_experts) + + check_case(16, 4096, 128, 'identity', with_bias=False, with_physical_map=False) + check_case(129, 4352, 260, 'sigmoid', with_image_bias=True, with_mask=True, with_unmapped_topk_idx=True) + + inputs = make_inputs(128, 4096, 384, with_image_bias=True) + inputs['bias'].fill_(1e8) + inputs['image_bias'].fill_(-1e8) + check_correctness(inputs, make_outputs(inputs), 'sqrtsoftplus') + + check_routing_controls() + + num_topk = 8 + inputs = make_inputs(8, 4096, 128, with_bias=False, with_physical_map=False) + inputs['x'].zero_() + inputs['weight'].zero_() + topk_idx, _ = run_mega_gate(inputs, make_outputs(inputs, num_topk=num_topk), 'identity', num_topk=num_topk) + assert torch.equal(topk_idx, torch.arange(num_topk, dtype=torch.int64, device='cuda').expand(8, -1)) + + inputs = make_inputs(16, 4096, 128, with_physical_map=False) + outputs = make_outputs(inputs, with_unmapped_topk_idx=True, num_topk=num_topk) + unmapped_storage = torch.empty((16, num_topk + 3), dtype=torch.int64, device='cuda') + outputs['unmapped_topk_idx'] = unmapped_storage[:, :num_topk] + assert outputs['unmapped_topk_idx'].stride() == (num_topk + 3, 1) + check_correctness(inputs, outputs, 'sqrtsoftplus', num_topk=num_topk) + + inputs = make_inputs(32, 4096, 128, with_physical_map=False) + outputs = make_outputs(inputs, with_unmapped_topk_idx=True, num_topk=num_topk) + force_random = torch.ones((32,), dtype=torch.bool, device='cuda') + topk_idx, topk_weights = run_mega_gate( + inputs, outputs, 'sqrtsoftplus', num_topk=num_topk, force_random=force_random) + assert torch.all((topk_idx >= 0) & (topk_idx < 128)).item() + assert torch.all(topk_weights > 0).item() + sorted_topk_idx = topk_idx.sort(dim=1).values + assert torch.all(sorted_topk_idx[:, 1:] != sorted_topk_idx[:, :-1]).item() + assert torch.all(outputs['unmapped_topk_idx'] == -1).item() + + inputs = make_inputs(0, 4096, 128, with_physical_map=False) + topk_idx, topk_weights = run_mega_gate(inputs, make_outputs(inputs, num_topk=num_topk), + 'sqrtsoftplus', num_topk=num_topk) + assert topk_idx.shape == (0, num_topk) and topk_weights.shape == (0, num_topk) + + inputs = make_inputs(16, 4096, 256) + outputs = make_outputs(inputs, with_unmapped_topk_idx=True) + deep_gemm.use_deterministic_algorithms(True) + try: + assert deep_gemm.get_bf16_mega_gate_config(16, 4096, 256, NUM_TOPK)['num_split_k'] == 1 + run_mega_gate(inputs, outputs, 'sqrtsoftplus') + expected = {name: tensor.clone() for name, tensor in outputs.items()} + for _ in range(100): + run_mega_gate(inputs, outputs, 'sqrtsoftplus') + for name, tensor in outputs.items(): + assert_bitwise_equal(tensor, expected[name], f'Deterministic invocation: {name}') + finally: + deep_gemm.use_deterministic_algorithms(False) + + +if __name__ == '__main__': + torch.manual_seed(0) + test_mega_gate_api_contract() + test_mega_gate() diff --git a/tests/test_mega_mhc.py b/tests/test_mega_mhc.py new file mode 100644 index 0000000000..faeabaa73a --- /dev/null +++ b/tests/test_mega_mhc.py @@ -0,0 +1,266 @@ +import os +import sys + +import torch + +import deep_gemm +from deep_gemm.testing import ( + assert_bitwise_equal, + bench_kineto, + calc_diff, + count_bytes, + get_arch_major, + test_filter, +) +from deep_gemm.utils import align + +sys.path.append(os.path.join(os.path.dirname(__file__), '..', 'third-party')) +from tilelang_ops import ref_mhc # noqa: E402 + + +def enumerate_cases(): + for sf_layout in ('col', 'extra'): + for hidden in (4096, 7168): + for num_tokens in (1, 64, 65, 200, 1025, 4096, 32768): + yield sf_layout, num_tokens, hidden + + +def make_inputs(num_tokens: int, hidden: int, seed: int, is_shifted: bool): + hc_mult = 4 + num_hc_outputs = hc_mult * (hc_mult + 2) + generator = torch.Generator(device='cuda').manual_seed(seed) + + def randn(shape, dtype): + return torch.randn(shape, dtype=dtype, device='cuda', generator=generator) + + inputs = dict( + x=randn((num_tokens, hidden), torch.bfloat16), + residual=randn((num_tokens, hc_mult, hidden), torch.bfloat16), + post_mix=randn((num_tokens, hc_mult, 1), torch.float).sigmoid(), + comb_res_mix=torch.rand((num_tokens, hc_mult, hc_mult), device='cuda', generator=generator), + shifted_prev_mix=randn((num_tokens, hc_mult, 1), torch.float).sigmoid() if is_shifted else None, + fn=randn((num_hc_outputs, hc_mult * hidden), torch.float).mul_(1e-2), + mix_scales=randn((3,), torch.float).mul_(0.1), + mix_bases=randn((num_hc_outputs,), torch.float).mul_(0.1), + rmsnorm_weight=randn((hidden,), torch.float).mul_(0.1).add_(1).bfloat16(), + hc_mult=hc_mult, + hc_norm_eps=2e-5, + hc_pre_eps=3e-4, + hc_post_scale=1.25, + sinkhorn_eps=2e-6, + num_sinkhorn_iters=10, + rmsnorm_eps=7e-6, + rmsnorm_scale=1.25, + ) + for _ in range(3): + inputs['comb_res_mix'] /= inputs['comb_res_mix'].sum(-1, keepdim=True) + inputs['comb_res_mix'] /= inputs['comb_res_mix'].sum(-2, keepdim=True) + duplicates = tuple(i for i in (63, 64, 511, 512, num_tokens - 1) if 0 < i < num_tokens) + for name in ('x', 'residual', 'post_mix', 'comb_res_mix', 'shifted_prev_mix'): + if inputs[name] is not None: + inputs[name][list(set(duplicates))] = inputs[name][0].clone() + return inputs, duplicates + + +def make_outputs(inputs, sf_layout: str, shared_sf_block_m: int = 224): + num_tokens, hidden = inputs['x'].shape + is_shifted = inputs['shifted_prev_mix'] is not None + outputs = dict( + new_residual=torch.empty_like(inputs['residual']), + new_prev_mix=torch.empty_like(inputs['shifted_prev_mix']) if is_shifted else None, + new_post_mix=torch.empty_like(inputs['post_mix']), + new_comb_res_mix=torch.empty_like(inputs['comb_res_mix']), + y_bf16=torch.empty_like(inputs['x']), + ) + if sf_layout == 'bf16': + return outputs + sf_shape = (num_tokens, hidden // 128) + outputs['y_fp8'] = torch.empty((num_tokens, hidden), dtype=torch.float8_e4m3fn, device='cuda') + if sf_layout == 'col': + outputs['y_gemm_sf'] = torch.empty_strided( + sf_shape, (1, align(num_tokens, 4)), dtype=torch.int32, device='cuda') + else: + outputs['y_routed_sf'] = torch.empty(sf_shape, dtype=torch.int32, device='cuda') + num_rows = (num_tokens + shared_sf_block_m - 1) // shared_sf_block_m * align(shared_sf_block_m, 128) + storage = torch.empty_strided( + (num_rows, hidden // 128), (1, num_rows), dtype=torch.int32, device='cuda') + outputs.update( + y_shared_sf=storage[:num_tokens], + y_shared_sf_storage=storage, + shared_sf_block_m=shared_sf_block_m, + ) + return outputs + + +def run_mega_mhc(inputs, outputs, **overrides): + kwargs = {**inputs, **outputs} + kwargs.pop('y_shared_sf_storage', None) + kwargs.update(overrides) + deep_gemm.mega_mhc(**kwargs) + + +def logical_outputs(result, shared_sf_block_m=None): + names = ('new_residual', 'new_post_mix', 'new_comb_res_mix', 'y_bf16') + names += ('new_prev_mix',) if result.get('new_prev_mix') is not None else () + names += tuple(name for name in ('y_fp8', 'y_gemm_sf', 'y_routed_sf') if name in result) + values = {name: result[name] for name in names} + if 'y_shared_sf_storage' in result: + y_bf16 = result['y_bf16'] + rows = ref_mhc.extra_sf_rows(y_bf16, result.get('shared_sf_block_m', shared_sf_block_m)) + values['y_shared_sf'] = result['y_shared_sf_storage'][rows] + return values + + +def comparable(tensor: torch.Tensor): + return tensor.contiguous().view(torch.uint8) if tensor.dtype == torch.int32 else tensor + + +def check_correctness(actual, pytorch_ref, baseline_ref, case): + max_diffs = [0.0, 0.0, 0.0] + for name in actual: + tensors = tuple(comparable(result[name]) for result in (actual, pytorch_ref, baseline_ref)) + diffs = tuple(float(calc_diff(tensors[i], tensors[j])) for i, j in ((0, 1), (0, 2), (2, 1))) + pytorch_limit = 2e-4 if name == 'y_fp8' else 5e-5 + baseline_limit = 1e-5 if case[0] == 'normal' and name == 'y_fp8' else 1e-6 + limits = (pytorch_limit, baseline_limit, pytorch_limit) + assert all(diff < limit for diff, limit in zip(diffs, limits)), \ + f'{case}, {name=}, {diffs=}' + max_diffs = [max(old, new) for old, new in zip(max_diffs, diffs)] + return max_diffs + + +@test_filter(lambda: get_arch_major() == 10) +@torch.no_grad() +def test_mega_mhc_api_contract() -> None: + num_tokens, hidden = 64, 4096 + inputs, _ = make_inputs(num_tokens, hidden, 2026, False) + + # Deterministic mode selects fixed Split-K independently for every invocation + outputs = make_outputs(inputs, 'bf16') + deep_gemm.use_deterministic_algorithms(True) + try: + run_mega_mhc(inputs, outputs) + expected = {name: tensor.clone() for name, tensor in logical_outputs(outputs).items()} + run_mega_mhc(inputs, outputs) + for name, tensor in logical_outputs(outputs).items(): + assert_bitwise_equal(tensor, expected[name], f'Deterministic invocation: {name}') + finally: + deep_gemm.use_deterministic_algorithms(False) + + # A stream must initialize its split barriers before CUDA Graph capture + expected_outputs = make_outputs(inputs, 'bf16') + run_mega_mhc(inputs, expected_outputs) + torch.cuda.synchronize() + capture_stream = torch.cuda.Stream() + graph_outputs = make_outputs(inputs, 'bf16') + try: + with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=capture_stream): + graph_outputs['y_bf16'].zero_() + run_mega_mhc(inputs, graph_outputs) + except RuntimeError as error: + assert 'CaptureStatus::None' in str(error) + else: + raise AssertionError('Mega mHC must be warmed up before CUDA Graph capture') + + with torch.cuda.stream(capture_stream): + run_mega_mhc(inputs, graph_outputs) + capture_stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + run_mega_mhc(inputs, graph_outputs) + graph.replay() + capture_stream.synchronize() + for name, tensor in logical_outputs(graph_outputs).items(): + assert_bitwise_equal(tensor, expected_outputs[name], f'CUDA Graph: {name}') + + # FP8-only calls remain bitwise stable after a different-sized invocation + for is_shifted in (False, True): + inputs, _ = make_inputs(num_tokens, hidden, 2026, is_shifted) + small_inputs, _ = make_inputs(1, hidden, 2027, is_shifted) + outputs = make_outputs(inputs, 'extra') + run_mega_mhc(inputs, outputs) + names = ('y_fp8', 'y_routed_sf', 'y_shared_sf_storage') + expected = {name: outputs[name].clone() for name in names} + run_mega_mhc(small_inputs, make_outputs(small_inputs, 'extra')) + run_mega_mhc(inputs, outputs, y_bf16=None) + for name in names: + assert_bitwise_equal(outputs[name], expected[name], f'FP8-only: {name}') + + # Concurrent kernels on different streams use independent split barriers + inputs, _ = make_inputs(num_tokens, hidden, 2028, False) + expected_outputs = make_outputs(inputs, 'bf16') + run_mega_mhc(inputs, expected_outputs) + torch.cuda.synchronize() + + streams = (torch.cuda.Stream(), torch.cuda.Stream()) + stream_outputs = [] + for stream in streams: + with torch.cuda.stream(stream): + outputs = make_outputs(inputs, 'bf16') + run_mega_mhc(inputs, outputs) + stream_outputs.append(outputs) + for stream in streams: + stream.synchronize() + for stream_idx, outputs in enumerate(stream_outputs): + for name, tensor in logical_outputs(outputs).items(): + assert_bitwise_equal(tensor, expected_outputs[name], f'Multi-stream {stream_idx}: {name}') + + +@test_filter(lambda: get_arch_major() == 10) +@torch.no_grad() +def test_mega_mhc() -> None: + if not ref_mhc.has_baseline(): + return + for is_shifted in (True, False): + mode = 'shifted' if is_shifted else 'normal' + baseline_kernels = ref_mhc.SHIFTED_KERNELS if is_shifted else ref_mhc.NORMAL_KERNELS + print(f'Testing Mega {mode} mHC:') + for case_idx, (sf_layout, num_tokens, hidden) in enumerate(enumerate_cases()): + inputs, duplicates = make_inputs( + num_tokens, hidden, 2026 + case_idx + (0 if is_shifted else 1000), is_shifted) + outputs = make_outputs(inputs, sf_layout) + ref_kwargs = dict( + **inputs, sf_layout=sf_layout, + shared_sf_block_m=outputs.get('shared_sf_block_m', 0)) + + def run_baseline(): + return ref_mhc.mhc_baseline(**ref_kwargs) + + pytorch_ref, baseline_ref = ref_mhc.mhc_reference(**ref_kwargs), run_baseline() + run_mega_mhc(inputs, outputs) + actual = logical_outputs(outputs) + refs = tuple(logical_outputs(ref, outputs.get('shared_sf_block_m')) + for ref in (pytorch_ref, baseline_ref)) + diffs = check_correctness(actual, *refs, (mode, sf_layout, num_tokens, hidden)) + + expected = {name: tensor.clone() for name, tensor in actual.items()} + for _ in range(30): + run_mega_mhc(inputs, outputs) + actual = logical_outputs(outputs) + for name in actual: + assert_bitwise_equal(actual[name], expected[name], name) + for token_idx in duplicates: + for name, tensor in actual.items(): + assert_bitwise_equal(tensor[token_idx], tensor[0], f'{name}, token={token_idx}') + + kernel_t = bench_kineto(lambda: run_mega_mhc(inputs, outputs), + 'sm100_mega_mhc', suppress_kineto_output=True) + baseline_times = bench_kineto(run_baseline, baseline_kernels, suppress_kineto_output=True) + baseline_t = sum(baseline_times) + logical_inputs = [tensor for tensor in inputs.values() if isinstance(tensor, torch.Tensor)] + logical_bytes = count_bytes(*logical_inputs, *actual.values()) + # Materialized Norm/Cast boundary, plus the normal-only Res reread + intermediate_io_bytes = 2 * count_bytes(outputs['y_bf16']) + if not is_shifted: + intermediate_io_bytes += count_bytes(outputs['new_residual']) + print(f' > {sf_layout:5} T={num_tokens:5}, H={hidden:4}: {kernel_t * 1e6:6.1f} us, ' + f'{logical_bytes / kernel_t / 1e9:4.0f} GB/s ' + f'({(logical_bytes + intermediate_io_bytes) / kernel_t / 1e9:4.0f} incl. Scratch I/O) | ' + f'baseline {baseline_t * 1e6:6.1f} us, {baseline_t / kernel_t:5.2f}x | ' + f'diff K/P={diffs[0]:.1e}, K/B={diffs[1]:.1e}, B/P={diffs[2]:.1e}') + print() + + +if __name__ == '__main__': + test_mega_mhc_api_contract() + test_mega_mhc() diff --git a/tests/test_mega_moe.py b/tests/test_mega_moe.py index 9987e024a0..6223e2537e 100644 --- a/tests/test_mega_moe.py +++ b/tests/test_mega_moe.py @@ -98,7 +98,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace): mma_type=args.mma_type ) - # Cast weights into FP4 + # Cast routed weights into FP4 or FP8 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) @@ -108,6 +108,15 @@ def _cast_weights_to_fp4(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torc w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups) return w, w_sf + def _cast_weights_to_fp8(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + num_groups, n, k = bf16_weights.shape + w = torch.empty((num_groups, n, k), device='cuda', dtype=torch.float8_e4m3fn) + 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_fp8(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(): @@ -142,7 +151,7 @@ def create_inputs(): shared_l1_weights = shared_l2_weights = None if not is_bf16xbf16: - # FP8 path: cast inputs to FP8/FP4 with per-32 UE8M0 SF + # FP8 path: cast inputs and weights with per-32 UE8M0 SF assert hidden % 128 == 0 and intermediate_hidden % 128 == 0 and shared_intermediate_hidden % 128 == 0 block_m = deep_gemm.get_block_m_for_mega_moe( num_ranks, num_experts, buffer.num_max_tokens_per_rank, num_tokens, num_topk, args.mma_type) @@ -151,8 +160,9 @@ def create_inputs(): shared_x = (x_fp8, x_sf_tma) if num_shared_experts > 0: shared_l1_x_sf = _to_shared_mega_moe_sf_layout(x_sf, block_m, buffer.shared_l1_acts_sf.shape[0]) - l1_weights = _cast_weights_to_fp4(l1_weights) - l2_weights = _cast_weights_to_fp4(l2_weights) + cast_weights = _cast_weights_to_fp8 if args.mma_type == 'fp8xfp8' else _cast_weights_to_fp4 + l1_weights = cast_weights(l1_weights) + l2_weights = cast_weights(l2_weights) if num_shared_experts > 0: shared_l1_weights = _cast_fp8_for_mega_moe(shared_l1_weights)[0::2] shared_l2_weights = _cast_fp8_for_mega_moe(shared_l2_weights)[0::2] @@ -352,22 +362,24 @@ def run_baseline(): # HBM bytes: weights + activations + output num_touched_experts = torch.unique(gathered_topk_idx[gathered_topk_idx >= 0]).numel() - act_elem_size, weight_elem_size = (2, 2) if is_bf16xbf16 else (1, 0.5) + act_elem_size = 2 if is_bf16xbf16 else 1 + routed_weight_elem_size = 2 if is_bf16xbf16 else (1 if args.mma_type == 'fp8xfp8' else 0.5) + shared_weight_elem_size = 2 if is_bf16xbf16 else 1 num_routed_hbm_bytes = ( - num_touched_experts * intermediate_hidden * 2 * hidden * weight_elem_size # L1 weights - + num_touched_experts * hidden * intermediate_hidden * weight_elem_size # L2 weights - + num_recv_tokens * hidden * act_elem_size # L1 acts read - + num_recv_tokens * intermediate_hidden * act_elem_size # L1 output write - + num_recv_tokens * intermediate_hidden * act_elem_size # L2 acts read - + num_recv_tokens * hidden * 2 # L2 output write (always BF16) + num_touched_experts * intermediate_hidden * 2 * hidden * routed_weight_elem_size # L1 weights + + num_touched_experts * hidden * intermediate_hidden * routed_weight_elem_size # L2 weights + + num_recv_tokens * hidden * act_elem_size # L1 acts read + + num_recv_tokens * intermediate_hidden * act_elem_size # L1 output write + + num_recv_tokens * intermediate_hidden * act_elem_size # L2 acts read + + num_recv_tokens * hidden * 2 # L2 output write (always BF16) ) num_shared_hbm_bytes = 0 if num_shared_experts == 0 else ( - shared_intermediate_hidden * 2 * hidden * weight_elem_size # Shared L1 weights - + hidden * shared_intermediate_hidden * weight_elem_size # Shared L2 weights - + num_tokens * hidden * act_elem_size # Shared L1 acts read - + num_tokens * shared_intermediate_hidden * act_elem_size # Shared L1 output write - + num_tokens * shared_intermediate_hidden * act_elem_size # Shared L2 acts read - + num_tokens * hidden * 2 # Shared L2 output write + shared_intermediate_hidden * 2 * hidden * shared_weight_elem_size # Shared L1 weights + + hidden * shared_intermediate_hidden * shared_weight_elem_size # Shared L2 weights + + num_tokens * hidden * act_elem_size # Shared L1 acts read + + num_tokens * shared_intermediate_hidden * act_elem_size # Shared L1 output write + + num_tokens * shared_intermediate_hidden * act_elem_size # Shared L2 acts read + + num_tokens * hidden * 2 # Shared L2 output write ) num_hbm_bytes = num_routed_hbm_bytes + num_shared_hbm_bytes @@ -422,7 +434,7 @@ def print_perf(elapsed: float, ref_time: float, ref_label: str): 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)') - parser.add_argument('--mma-type', type=str, default='fp8xfp4', help='MMA type: fp8xfp4 or bf16xbf16') + parser.add_argument('--mma-type', type=str, default='fp8xfp4', choices=('fp8xfp4', 'fp8xfp8', 'bf16xbf16')) # Test settings parser.add_argument('--num-correctness-tests', type=int, default=None, help='Pressure test') diff --git a/tests/test_sanitizer.py b/tests/test_sanitizer.py index 75ab10e60a..00bfb60fa3 100644 --- a/tests/test_sanitizer.py +++ b/tests/test_sanitizer.py @@ -51,7 +51,7 @@ env = os.environ.copy() env['CUDA_LAUNCH_BLOCKING'] = '1' - env['DG_JIT_PTXAS_CHECK'] = '1' + env['DG_JIT_CHECK_NO_SPILLS'] = '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' diff --git a/tests/utils.py b/tests/utils.py index 143f78a6b7..ed40a2a2c9 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,6 +1,9 @@ import torch +from typing import Callable, Tuple -from deep_gemm.utils import align, get_mk_alignment_for_contiguous_layout +import deep_gemm +from deep_gemm.testing import assert_bitwise_equal +from deep_gemm.utils import align, ceil_div, get_mk_alignment_for_contiguous_layout def assert_psum_zero_padding(a: torch.Tensor | tuple, d: torch.Tensor, grouped_layout: torch.Tensor, dtype_label: str) -> None: @@ -12,3 +15,65 @@ def assert_psum_zero_padding(a: torch.Tensor | tuple, d: torch.Tensor, grouped_l d_padding = d[current_m: aligned_m] assert torch.equal(a_padding, torch.zeros_like(a_padding)), f'{group_idx=}, nonzero {dtype_label} input padding' assert torch.equal(d_padding, torch.zeros_like(d_padding)), f'{group_idx=}, nonzero {dtype_label} output padding' + + +def assert_direct_output_matches_fp32_accumulation( + direct_output: torch.Tensor, + launch: Callable[[torch.Tensor, torch.Tensor | None], None], + label: str) -> None: + """Check direct output against the same workload accumulated from zero in FP32. + + ``launch(output, accumulator)`` reruns the direct-output workload with the + same inputs, shapes, group layout, K lengths, and kernel options, changing + only the output mode. This helper passes one zero-initialized FP32 tensor as + both D and C, so accumulation mode computes ``D = GEMM + 0``. FP32 direct + output must match this baseline bitwise; BF16 direct output must match the + baseline after one FP32-to-BF16 cast. + """ + assert direct_output.dtype in (torch.float, torch.bfloat16) + + fp32_baseline = torch.zeros_like(direct_output, dtype=torch.float) + launch(fp32_baseline, fp32_baseline) + expected = fp32_baseline if direct_output.dtype == torch.float else fp32_baseline.to(torch.bfloat16) + assert_bitwise_equal(direct_output, expected, label) + + +def convert_to_fp8(x: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]: + # E4M3 exactly represents every finite E2M1 value. Convert the encoding without requantizing. + data, sf = x + if data.dtype == torch.float8_e4m3fn: + return x + if data.stride(-1) != 1: + converted, sf = convert_to_fp8((data.mT, sf)) + return converted.mT, sf + lut = torch.tensor((0x00, 0x30, 0x38, 0x3c, 0x40, 0x44, 0x48, 0x4c), dtype=torch.uint8, device=data.device) + packed = data.view(torch.uint8) + lo, hi = packed & 0xf, packed >> 4 + codes = torch.stack((lut[(lo & 7).long()] | ((lo & 8) << 4), + lut[(hi & 7).long()] | ((hi & 8) << 4)), dim=-1).flatten(-2) + return codes.view(torch.float8_e4m3fn), sf + + +def to_cublaslt_vec16_sf_layout(sf: torch.Tensor) -> torch.Tensor: + # Convert per-32 MXFP4 power-of-two scales into the equivalent per-16 NVFP4 layout. + assert sf.dim() == 2 and sf.dtype == torch.float + assert (sf.view(torch.int32) & 0x7fffff == 0).all(), 'SF must be positive powers of two' + exp = (sf.view(torch.int32) >> 23) - 127 + assert (exp >= -9).all() and (exp <= 8).all(), 'SF out of the UE4M3 range' + ue4m3 = torch.where(exp >= -6, (exp + 7) << 3, 1 << (exp + 9)).to(torch.uint8).repeat_interleave(2, dim=1) + + mn, sf_k = ue4m3.shape + num_mn_tiles, num_k_tiles = ceil_div(mn, 128), ceil_div(sf_k, 4) + padded = torch.zeros((num_mn_tiles * 128, num_k_tiles * 4), device=sf.device, dtype=torch.uint8) + padded[:mn, :sf_k] = ue4m3 + return padded.view(num_mn_tiles, 4, 32, num_k_tiles, 4).permute(0, 3, 2, 1, 4).contiguous() + + +def make_cublas_gemm(a: Tuple[torch.Tensor, torch.Tensor], b: Tuple[torch.Tensor, torch.Tensor], + d: torch.Tensor, c: torch.Tensor | None) -> Callable[[], None]: + if a[0].dtype == torch.int8 and b[0].dtype == torch.int8: + a_nvfp4 = a[0], to_cublaslt_vec16_sf_layout(a[1]) + b_nvfp4 = b[0], to_cublaslt_vec16_sf_layout(b[1]) + return lambda: deep_gemm.cublaslt_nvfp4_gemm_nt(a_nvfp4, b_nvfp4, d, c=c) + a_fp8, b_fp8 = convert_to_fp8(a), convert_to_fp8(b) + return lambda: deep_gemm.cublaslt_gemm_nt(a_fp8[0], b_fp8[0], d, c=c) diff --git a/third-party/deep_jit b/third-party/deep_jit new file mode 160000 index 0000000000..e5bdee2bc4 --- /dev/null +++ b/third-party/deep_jit @@ -0,0 +1 @@ +Subproject commit e5bdee2bc4ca519eba00cfc5f0c6e950e6a96a16 diff --git a/third-party/fmt b/third-party/fmt deleted file mode 160000 index 553ec11ec0..0000000000 --- a/third-party/fmt +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 553ec11ec06fbe0beebfbb45f9dc3c9eabd83d28 diff --git a/third-party/tilelang_ops/norm.py b/third-party/tilelang_ops/norm.py new file mode 100644 index 0000000000..b4a7c57195 --- /dev/null +++ b/third-party/tilelang_ops/norm.py @@ -0,0 +1,213 @@ +# Adapted from vLLM's TileLang norm kernel for the Mega mHC baseline. +import tilelang +import tilelang.language as T +import torch + + +FP8_AMAX_MARGIN = 1e-4 +FLOAT8E4M3_MAX = 448.0 + + +@T.macro +def ue8m0_exp_scale(scale: T.float32) -> T.int32: + scale_bits = T.reinterpret("uint32", scale) + exponent = (scale_bits >> 23) & 0xFF + has_mantissa = scale_bits & ((1 << 23) - 1) != 0 + return -T.reinterpret("int32", exponent - 127 + has_mantissa) + + +@tilelang.jit( + out_idx=None, + target="cuda", + pass_configs={"tl.disable_tma_lower": True}, +) +def norm_tl( + hidden: int, + eps: float, + out_scale: float, + sf_layout: str, + shared_sf_block_m: int, +): + assert hidden in (4096, 7168) + assert sf_layout in ("bf16", "col", "extra") + + do_fp8_cast = sf_layout != "bf16" + num_tokens = T.dynamic("num_tokens") + block_hidden = tilelang.next_power_of_2(hidden) + threads = 256 + vec_size = 8 if hidden % 2048 == 0 else 4 + values_per_thread = block_hidden // threads + + def local_layout(i: int, j: int) -> tuple[int, int]: + thread_idx = j % (threads * vec_size) // vec_size + local_idx = i * values_per_thread + j // (threads * vec_size) * vec_size + j % vec_size + return thread_idx, local_idx + + num_scales = hidden // 32 + block_num_scales = block_hidden // 32 + output_dtype = T.float8_e4m3fn if do_fp8_cast else "bfloat16" + if sf_layout == "col": + num_scale_packs = num_scales // 4 + aligned_tokens = (num_tokens + 3) // 4 * 4 + sf_type = "uint8" + sf_shape = (num_tokens, num_scale_packs, 4) + sf_stride = (4, aligned_tokens * 4, 1) + elif sf_layout == "extra": + sf_type = "uint8" + sf_shape = (num_tokens, num_scales) + sf_stride = (num_scales, 1) + else: + sf_type = "bfloat16" + sf_shape = (num_tokens, hidden) + sf_stride = (hidden, 1) + + if sf_layout == "extra": + aligned_shared_block_m = (shared_sf_block_m + 127) // 128 * 128 + shared_sf_rows = T.dynamic("shared_sf_rows") + shared_sf_type = "uint8" + shared_sf_shape = (shared_sf_rows, hidden // 128, 4) + shared_sf_stride = (4, shared_sf_rows * 4, 1) + + def shared_sf_row(token_idx): + index_in_block = token_idx % shared_sf_block_m + return ( + token_idx // shared_sf_block_m * aligned_shared_block_m + + index_in_block // 128 * 128 + + index_in_block % 32 * 4 + + index_in_block % 128 // 32 + ) + else: + shared_sf_type = "bfloat16" + shared_sf_shape = (num_tokens, hidden) + shared_sf_stride = (hidden, 1) + + @T.prim_func + def _norm( + x: T.Tensor((num_tokens, hidden), "bfloat16"), # type: ignore + y: T.Tensor((num_tokens, hidden), output_dtype), # type: ignore + weight: T.Tensor((hidden,), "bfloat16"), # type: ignore + sf: T.StridedTensor(sf_shape, sf_stride, sf_type), # type: ignore + shared_sf: T.StridedTensor(shared_sf_shape, shared_sf_stride, shared_sf_type), # type: ignore + y_bf16: T.Tensor((num_tokens, hidden), "bfloat16"), # type: ignore + ) -> None: + with T.Kernel(num_tokens, threads=threads) as (token_idx): + x_local = T.alloc_fragment((1, block_hidden), "float") + y_local = T.alloc_fragment((1, block_hidden), "float") + square_local = T.alloc_fragment((1, block_hidden), "float") + rstd_local = T.alloc_fragment((1,), "float") + weight_local = T.alloc_fragment((block_hidden,), "bfloat16") + + if do_fp8_cast: + y_quant_local = T.alloc_fragment((1, block_hidden), "float") + y_fp8_local = T.alloc_fragment((1, block_hidden), T.float8_e4m3fn) + amax_local = T.alloc_fragment((1, block_num_scales), "float") + scale_local = T.alloc_fragment((1, block_num_scales), "float") + sf_local = T.alloc_fragment((1, block_num_scales), "uint8") + T.annotate_layout({ + x_local: T.Fragment(x_local.shape, forward_fn=local_layout), + y_quant_local: T.Fragment(y_quant_local.shape, forward_fn=local_layout), + y_fp8_local: T.Fragment(y_fp8_local.shape, forward_fn=local_layout), + }) + else: + T.annotate_layout({ + x_local: T.Fragment(x_local.shape, forward_fn=local_layout), + }) + + T.copy(weight, weight_local[:hidden]) + T.pdl_sync() + T.copy(x[token_idx, :], x_local[0, :hidden]) + for j in T.Parallel(block_hidden): + square_local[0, j] = x_local[0, j] * x_local[0, j] if j < hidden else 0 + T.reduce_sum(square_local, rstd_local, dim=1) + rstd_local[0] = T.rsqrt(rstd_local[0] / hidden + eps) + for j in T.Parallel(block_hidden): + if j < hidden: + y_local[0, j] = x_local[0, j] * rstd_local[0] * (weight_local[j] * out_scale) + + if not do_fp8_cast: + T.pdl_trigger() + T.copy(y_local[0, :hidden], y[token_idx, :]) + else: + T.copy(y_local[0, :hidden], y_bf16[token_idx, :]) + for j in T.Parallel(block_hidden): + y_quant_local[0, j] = T.cast(T.cast(y_local[0, j], "bfloat16"), "float") + y_quant_reshaped = T.reshape(y_quant_local, [1, block_num_scales, 32]) + T.reduce_absmax(y_quant_reshaped, amax_local, dim=2) + + for j in T.Parallel(block_num_scales): + amax_local[0, j] = T.max(amax_local[0, j], FP8_AMAX_MARGIN) + exponent = ue8m0_exp_scale(amax_local[0, j] / FLOAT8E4M3_MAX) + scale_local[0, j] = T.reinterpret("float32", (exponent + 127) << 23) + sf_local[0, j] = T.cast(-exponent + 127, "uint8") + + for j in T.Parallel(block_hidden): + y_fp8_local[0, j] = T.cast( + y_quant_local[0, j] * scale_local[0, j // 32], T.float8_e4m3fn) + y[token_idx, j] = y_fp8_local[0, j] + + T.pdl_trigger() + if sf_layout == "col": + for j in T.Parallel(block_num_scales): + if j < num_scales: + sf[token_idx, j // 4, j % 4] = sf_local[0, j] + else: + T.copy(sf_local[0, :num_scales], sf[token_idx, :]) + for j in T.Parallel(block_num_scales): + if j < num_scales: + shared_sf[shared_sf_row(token_idx), j // 4, j % 4] = sf_local[0, j] + + return _norm + + +def norm( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + out_scale: float, + sf_layout: str = "bf16", + shared_sf_out: torch.Tensor | None = None, + shared_sf_block_m: int = 0, +): + assert sf_layout in ("bf16", "col", "extra") + assert x.ndim == 2 and x.dtype == torch.bfloat16 and x.is_contiguous() + num_tokens, hidden = x.shape + assert hidden in (4096, 7168) + assert weight.shape == (hidden,) and weight.dtype == torch.bfloat16 and weight.is_contiguous() + + do_fp8_cast = sf_layout != "bf16" + y_bf16 = torch.empty(x.shape, dtype=torch.bfloat16, device=x.device) + if do_fp8_cast: + y = torch.empty(x.shape, dtype=torch.float8_e4m3fn, device=x.device) + num_scale_packs = hidden // 128 + if sf_layout == "col": + aligned_tokens = (num_tokens + 3) // 4 * 4 + sf_storage = torch.empty( + num_scale_packs, aligned_tokens, dtype=torch.int32, device=x.device) + sf = sf_storage.mT[:num_tokens] + sf_arg = sf_storage.view(torch.uint8).as_strided( + (num_tokens, num_scale_packs, 4), (4, aligned_tokens * 4, 1)) + else: + sf = torch.empty( + num_tokens, num_scale_packs, dtype=torch.int32, device=x.device) + sf_arg = sf.view(torch.uint8) + else: + y = y_bf16 + sf = None + sf_arg = x + + if sf_layout == "extra": + assert shared_sf_out is not None and shared_sf_block_m > 0 + num_rows = shared_sf_out.stride(1) + assert shared_sf_out.shape == (num_tokens, hidden // 128) + assert shared_sf_out.dtype == torch.int32 and shared_sf_out.stride(0) == 1 + aligned_block_m = (shared_sf_block_m + 127) // 128 * 128 + assert (num_tokens + shared_sf_block_m - 1) // shared_sf_block_m * aligned_block_m <= num_rows + shared_sf_arg = shared_sf_out.mT.view(torch.uint8).as_strided( + (num_rows, hidden // 128, 4), (4, num_rows * 4, 1)) + else: + assert shared_sf_out is None and shared_sf_block_m == 0 + shared_sf_arg = x + + kernel = norm_tl(hidden, eps, out_scale, sf_layout, shared_sf_block_m) + kernel(x, y, weight, sf_arg, shared_sf_arg, y_bf16) + return (y_bf16, y, sf) if do_fp8_cast else y diff --git a/third-party/tilelang_ops/ref_mhc.py b/third-party/tilelang_ops/ref_mhc.py new file mode 100644 index 0000000000..f084055799 --- /dev/null +++ b/third-party/tilelang_ops/ref_mhc.py @@ -0,0 +1,277 @@ +import torch + +import deep_gemm + +from deep_gemm.utils import align, per_token_cast_to_fp8 + + +_BASELINE_IMPORT_ERROR = None +try: + from tile_kernels.modeling.mhc.ops import ( + MHCGemmWithSqrsum, + mhc_post, + mhc_pre_apply_mix, + mhc_pre_big_fuse, + mhc_pre_split_mixes, + sinkhorn_normalize, + ) + from tile_kernels.modeling.mhc.ops.norm_fn import MHCReducePartialsAndRmsnorm + + from .norm import norm + + def deep_gemm_hc_gemm_with_sqr_sum( + x: torch.Tensor, fn: torch.Tensor, transpose_fn: bool, num_splits: int, + ): + assert not transpose_fn + x_flat = x.view(-1, x.size(-2) * x.size(-1)) + output_shape = (num_splits, x_flat.size(0)) + gemm_out = torch.empty( + (*output_shape, fn.size(0)), dtype=torch.float32, device=x.device) + sqr_sum = torch.empty(output_shape, dtype=torch.float32, device=x.device) + deep_gemm.tf32_hc_prenorm_gemm( + x_flat, fn, gemm_out, sqr_sum, num_splits=num_splits) + return gemm_out.unsqueeze(2), sqr_sum.unsqueeze(2) + + # Keep the fused Normal post-GEMM path, replacing only its GEMM+sqr-sum producer. + MHCGemmWithSqrsum.apply = staticmethod(deep_gemm_hc_gemm_with_sqr_sum) +except Exception as ex: + _BASELINE_IMPORT_ERROR = ex + + +NUM_SPLITS = 16 +SHIFTED_KERNELS = ( + 'mhc_post_fwd', + 'mhc_pre_apply_mix_fwd', + 'tf32_hc_prenorm_gemm', + 'mhc_reduce_partials_and_rmsnorm_fwd', + 'mhc_pre_split_mixes_fwd', + 'mhc_sinkhorn_fwd', + '_norm_kernel', +) +NORMAL_KERNELS = ( + 'mhc_post_fwd', + 'tf32_hc_prenorm_gemm', + 'mhc_pre_big_fuse_kernel', + '_norm_kernel', +) + + +def has_baseline() -> bool: + if _BASELINE_IMPORT_ERROR is None: + return True + print(f'Failed to load baseline code: {_BASELINE_IMPORT_ERROR}, skip baseline tests') + return False + + +def extra_sf_rows(x: torch.Tensor, block_m: int) -> torch.Tensor: + token_idx = torch.arange(x.size(0), device=x.device) + index_in_block = token_idx % block_m + return ( + token_idx // block_m * align(block_m, 128) + + index_in_block // 128 * 128 + + index_in_block % 32 * 4 + + index_in_block % 128 // 32 + ) + + +def _allocate_extra_sf(x: torch.Tensor, block_m: int) -> torch.Tensor: + num_tokens, hidden = x.shape + num_rows = (num_tokens + block_m - 1) // block_m * align(block_m, 128) + return torch.empty_strided( + (num_rows, hidden // 128), (1, num_rows), dtype=torch.int32, device=x.device) + + +def _norm_outputs( + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + rmsnorm_scale: float, + sf_layout: str, + shared_sf_block_m: int, +) -> dict[str, torch.Tensor]: + if sf_layout == 'bf16': + return {'y_bf16': norm(x, weight, eps, rmsnorm_scale)} + + extra_storage = _allocate_extra_sf(x, shared_sf_block_m) if sf_layout == 'extra' else None + y_bf16, y_fp8, y_fp8_sf = norm( + x, + weight, + eps, + rmsnorm_scale, + sf_layout, + extra_storage[:x.size(0)] if extra_storage is not None else None, + shared_sf_block_m, + ) + sf_name = 'y_gemm_sf' if sf_layout == 'col' else 'y_routed_sf' + result = {'y_bf16': y_bf16, 'y_fp8': y_fp8, sf_name: y_fp8_sf} + if extra_storage is not None: + result['y_shared_sf_storage'] = extra_storage + return result + + +@torch.no_grad() +def mhc_baseline( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + mix_scales: torch.Tensor, + mix_bases: torch.Tensor, + rmsnorm_weight: torch.Tensor, + hc_mult: int, + hc_norm_eps: float, + hc_pre_eps: float, + hc_post_scale: float, + sinkhorn_eps: float, + num_sinkhorn_iters: int, + rmsnorm_eps: float, + rmsnorm_scale: float, + sf_layout: str = 'bf16', + shared_sf_block_m: int = 0, + shifted_prev_mix: torch.Tensor | None = None, +) -> dict[str, torch.Tensor]: + new_residual = mhc_post( + x.unsqueeze(0), residual.unsqueeze(0), + post_mix.unsqueeze(0), comb_res_mix.unsqueeze(0), + ).squeeze(0) + + if shifted_prev_mix is None: + new_post_mix, new_comb_res_mix, norm_input = mhc_pre_big_fuse( + new_residual, fn, mix_scales, mix_bases, + hc_norm_eps, hc_pre_eps, sinkhorn_eps, hc_post_scale, + num_sinkhorn_iters, n_splits=NUM_SPLITS, + ) + result = { + 'new_residual': new_residual, + 'new_post_mix': new_post_mix, + 'new_comb_res_mix': new_comb_res_mix, + } + else: + norm_input = mhc_pre_apply_mix( + new_residual.unsqueeze(0), shifted_prev_mix.unsqueeze(0), + ).squeeze(0) + gemm_out, hc_sqr_sum = MHCGemmWithSqrsum.apply( + new_residual.unsqueeze(0), fn, False, NUM_SPLITS, + ) + mixes = MHCReducePartialsAndRmsnorm.apply( + gemm_out, hc_sqr_sum, hc_norm_eps, hc_mult * x.size(1), + ) + new_prev_mix, new_post_mix, new_comb_res_mix = mhc_pre_split_mixes( + mixes.unsqueeze(0), mix_scales, mix_bases, + hc_mult, hc_post_scale, hc_pre_eps, + ) + new_comb_res_mix = sinkhorn_normalize( + new_comb_res_mix, repeat=num_sinkhorn_iters, eps=sinkhorn_eps, + ) + result = { + 'new_residual': new_residual, + 'new_prev_mix': new_prev_mix.squeeze(0), + 'new_post_mix': new_post_mix.squeeze(0), + 'new_comb_res_mix': new_comb_res_mix.squeeze(0), + } + + result.update(_norm_outputs( + norm_input, rmsnorm_weight, rmsnorm_eps, rmsnorm_scale, + sf_layout, shared_sf_block_m, + )) + return result + + +def _sinkhorn_reference(value: torch.Tensor, repeat: int, eps: float) -> torch.Tensor: + value = value.softmax(dim=-1) + eps + value = value / (value.sum(dim=-2, keepdim=True) + eps) + for _ in range(repeat - 1): + value = value / (value.sum(dim=-1, keepdim=True) + eps) + value = value / (value.sum(dim=-2, keepdim=True) + eps) + return value + + +@torch.no_grad() +def mhc_reference( + x: torch.Tensor, + residual: torch.Tensor, + post_mix: torch.Tensor, + comb_res_mix: torch.Tensor, + fn: torch.Tensor, + mix_scales: torch.Tensor, + mix_bases: torch.Tensor, + rmsnorm_weight: torch.Tensor, + hc_mult: int, + hc_norm_eps: float, + hc_pre_eps: float, + hc_post_scale: float, + sinkhorn_eps: float, + num_sinkhorn_iters: int, + rmsnorm_eps: float, + rmsnorm_scale: float, + sf_layout: str = 'bf16', + shared_sf_block_m: int = 0, + shifted_prev_mix: torch.Tensor | None = None, +) -> dict[str, torch.Tensor]: + new_residual = x.float().unsqueeze(1) * post_mix + new_residual += torch.einsum( + 'tij,tih->tjh', comb_res_mix, residual.float()) + new_residual = new_residual.bfloat16() + norm_input = None if shifted_prev_mix is None else ( + new_residual.float() * shifted_prev_mix).sum(dim=1).bfloat16() + + old_allow_tf32 = torch.backends.cuda.matmul.allow_tf32 + torch.backends.cuda.matmul.allow_tf32 = False + try: + mixes = new_residual.flatten(1).float() @ fn.mT + finally: + torch.backends.cuda.matmul.allow_tf32 = old_allow_tf32 + hc_sqr_sum = new_residual.float().square().sum((1, 2)) + mixes *= torch.rsqrt( + hc_sqr_sum / (hc_mult * x.size(1)) + hc_norm_eps).unsqueeze(1) + scales = torch.cat(( + mix_scales[0].expand(hc_mult), + mix_scales[1].expand(hc_mult), + mix_scales[2].expand(hc_mult * hc_mult), + )) + mixes = mixes * scales + mix_bases + new_prev_mix = mixes[:, :hc_mult].sigmoid().unsqueeze(2) + hc_pre_eps + new_post_mix = ( + mixes[:, hc_mult:2 * hc_mult].sigmoid() * hc_post_scale + ).unsqueeze(2) + new_comb_res_mix = _sinkhorn_reference( + mixes[:, 2 * hc_mult:].view(-1, hc_mult, hc_mult), + num_sinkhorn_iters, + sinkhorn_eps, + ) + if norm_input is None: + norm_input = (new_residual.float() * new_prev_mix).sum(dim=1).bfloat16() + + norm_input_float = norm_input.float() + y_bf16 = ( + norm_input_float + * torch.rsqrt(norm_input_float.square().mean(1) + rmsnorm_eps).unsqueeze(1) + * (rmsnorm_weight.float() * rmsnorm_scale) + ).bfloat16() + result = { + 'new_residual': new_residual, + 'new_post_mix': new_post_mix, + 'new_comb_res_mix': new_comb_res_mix, + 'y_bf16': y_bf16, + } + if shifted_prev_mix is not None: + result['new_prev_mix'] = new_prev_mix + if sf_layout == 'bf16': + return result + + y_fp8, y_fp8_sf = per_token_cast_to_fp8( + y_bf16, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True) + if sf_layout == 'col': + col_major_sf = torch.empty_strided( + y_fp8_sf.shape, (1, align(x.size(0), 4)), + dtype=torch.int32, device=x.device) + col_major_sf.copy_(y_fp8_sf) + y_fp8_sf = col_major_sf + sf_name = 'y_gemm_sf' if sf_layout == 'col' else 'y_routed_sf' + result.update(y_fp8=y_fp8, **{sf_name: y_fp8_sf}) + if sf_layout == 'extra': + storage = _allocate_extra_sf(x, shared_sf_block_m) + storage[extra_sf_rows(x, shared_sf_block_m)] = y_fp8_sf + result['y_shared_sf_storage'] = storage + return result From 78b69000794d0937b47ae3387eff7663410264d1 Mon Sep 17 00:00:00 2001 From: Zhean Xu <94977922+zheanxu@users.noreply.github.com> Date: Mon, 14 Sep 2026 13:46:03 +0800 Subject: [PATCH 2/5] Fix Mega MoE task info slot release ordering (#441) --- deep_gemm/include/deep_gemm/ptx/ld_st.cuh | 5 +++++ deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh | 2 ++ 2 files changed, 7 insertions(+) diff --git a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh index 5cfec6aaab..a94a8a5e71 100644 --- a/deep_gemm/include/deep_gemm/ptx/ld_st.cuh +++ b/deep_gemm/include/deep_gemm/ptx/ld_st.cuh @@ -98,6 +98,11 @@ struct SM100_U8x8_STSM_T { } }; +/// Memory ordering +CUTLASS_DEVICE void fence_acq_rel_cta() { + asm volatile("fence.acq_rel.cta;" ::: "memory"); +} + /// Shared memory CUTLASS_DEVICE uint32_t ld_shared(const uint32_t* ptr) { uint32_t ret; diff --git a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh index c584bcf828..ccc17e7d9b 100644 --- a/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh @@ -246,6 +246,8 @@ struct MegaMoEScheduler { } CUTLASS_DEVICE void release_task_info() const { + // Complete metadata reads before the scheduler can overwrite this slot. + ptx::fence_acq_rel_cta(); task_info_empty_barriers[sched_stage_idx ^ 1].arrive(0u); } From 1364fed0f23ef7b5c99a10bd0b25f846af5848c7 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Tue, 22 Sep 2026 16:57:06 +0000 Subject: [PATCH 3/5] Reuse a warmed sparse MQA logits workspace during graph capture Frameworks such as TensorRT-LLM capture CUDA graphs on a dedicated side stream after warming up eagerly on the launching stream. The sparse MQA logits workspace is keyed by stream, so the first call on the capture stream used to trip the pre-capture assertion. Captured work replays on the launching stream, so fall back to that stream's warmed-up workspace while capturing instead of recording a new allocation into the graph. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- csrc/apis/attention.hpp | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 5b8e603d65..8bd9cf02e9 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -263,7 +263,16 @@ static const torch::Tensor& get_sparse_mqa_logits_workspace(const torch::TensorO auto& workspace = workspaces[stream]; if (not workspace.defined()) { // Warm up each stream before capture so one-time zeroing is not replayed with the graph. - DG_HOST_ASSERT(c10::cuda::currentStreamCaptureStatusMayInitCtx() == c10::cuda::CaptureStatus::None); + if (c10::cuda::currentStreamCaptureStatusMayInitCtx() != c10::cuda::CaptureStatus::None) { + // Frameworks capture on a dedicated side stream but replay on the stream that + // launched the eager warm-up, so reuse that stream's workspace instead of + // recording a fresh allocation and its zeroing into the graph. + for (const auto& [other_stream, other_workspace] : workspaces) { + if (other_workspace.defined() and other_workspace.device() == options.device()) + return other_workspace; + } + DG_HOST_UNREACHABLE("Warm up the sparse MQA logits workspace on this device before capturing"); + } workspace = torch::zeros({kNumWorkspaceBytes}, options.dtype(torch::kByte)); } return workspace; From 7a85767db7509d10e6348c00585179c414ab9f8b Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Tue, 22 Sep 2026 18:48:23 +0000 Subject: [PATCH 4/5] Return a typed map from get_bf16_mega_gate_config pybind11 converts std::map to a Python dict exactly like py::dict, but generated stubs become dict[str, int] instead of a bare dict, which strict type checkers reject. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- csrc/apis/mega_gate.hpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/csrc/apis/mega_gate.hpp b/csrc/apis/mega_gate.hpp index 617a43fb81..c02e57cf72 100644 --- a/csrc/apis/mega_gate.hpp +++ b/csrc/apis/mega_gate.hpp @@ -1,5 +1,8 @@ #pragma once +#include +#include + #include #include #include @@ -191,7 +194,7 @@ bf16_mega_gate(const torch::Tensor& x, return {topk_idx, topk_weights}; } -static pybind11::dict get_bf16_mega_gate_config(const int& num_tokens, const int& hidden, +static std::map get_bf16_mega_gate_config(const int& num_tokens, const int& hidden, const int& num_routed_experts, const int& num_topk) { DG_HOST_ASSERT(num_tokens > 0 and num_tokens <= static_cast(layout::mega_gate::kNumMaxTokens)); DG_HOST_ASSERT(hidden > 0 and hidden % 256 == 0); @@ -200,7 +203,8 @@ static pybind11::dict get_bf16_mega_gate_config(const int& num_tokens, const int const auto num_device_sms = runtime->get_num_sms(); const auto config = get_sm100_bf16_mega_gate_config(num_tokens, hidden, num_routed_experts, num_device_sms, true, true, true); - pybind11::dict result; + // A typed map keeps generated Python stubs precise (`dict[str, int]`) + std::map result; result["block_tokens"] = config.block_tokens; result["num_mma_ctas"] = config.num_mma_ctas; result["num_split_k"] = config.num_split_k; From d0a7e771f99fea7cc27513e46059de4f8761cc24 Mon Sep 17 00:00:00 2001 From: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> Date: Tue, 22 Sep 2026 20:24:54 +0000 Subject: [PATCH 5/5] Accept 16-byte aligned KV pages in the paged sparse MQA logits The paged sparse kernel addresses pages directly and copies them with 16-byte cp.async chunks, so a page stride only has to be a multiple of 16 bytes. Caches whose page-footer layout is not padded to 512 bytes (for example 64 x (64 + 4) = 4352-byte pages) can now be read in place. Signed-off-by: Fanrong Li <23290157+lfr-0531@users.noreply.github.com> --- csrc/apis/attention.hpp | 2 +- deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/csrc/apis/attention.hpp b/csrc/apis/attention.hpp index 8bd9cf02e9..28628b2590 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -415,7 +415,7 @@ static torch::Tensor fp8_fp4_paged_sparse_mqa_logits(const std::tuple(sizeof(int))); DG_HOST_ASSERT(fused_kv_cache.scalar_type() == torch::kUInt8 and fused_kv_cache.stride(1) == head_dim_with_sf and fused_kv_cache.stride(3) == 1 and fused_kv_cache.stride(0) <= std::numeric_limits::max() and - fused_kv_cache.stride(0) % 512 == 0); + fused_kv_cache.stride(0) % 16 == 0); // pages are read with 16-byte cp.async chunks const auto [_num_q_tokens_weights, _num_heads_weights] = get_shape<2>(weights); DG_HOST_ASSERT(_num_q_tokens_weights == num_q_tokens and _num_heads_weights == num_heads); diff --git a/deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh b/deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh index bfe36287b5..ee061fdc17 100644 --- a/deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh @@ -616,7 +616,7 @@ void sm100_paged_sparse_mqa_logits(const uint32_t logits_stride, const uint32_t // Keep the CUTLASS dtype out of the kernel template signature to avoid ptxas register spills using qk_dtype_t = cute::conditional_t; DG_STATIC_ASSERT(PAGE_KV % SPARSE_BLOCK_KV == 0, "Sparse KV blocks must not cross pages"); - DG_DEVICE_ASSERT(kv_page_stride_bytes % 512 == 0); + DG_DEVICE_ASSERT(kv_page_stride_bytes % 16 == 0); const auto kv_accessor = sparse_mqa_detail::PagedSparseKVAccessor( fused_kv_cache, kv_page_stride_bytes); sm100_sparse_mqa_logits_core_impl