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 73e6c710ee..330322eb81 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 @@ -176,6 +172,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 @@ -188,19 +185,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 @@ -212,9 +211,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 1abfd5c9b0..28628b2590 100644 --- a/csrc/apis/attention.hpp +++ b/csrc/apis/attention.hpp @@ -1,24 +1,27 @@ #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/sm120_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm100_mqa_logits.hpp" #include "../jit_kernels/impls/sm100_mqa_logits_f16_weights.hpp" +#include "../jit_kernels/impls/sm100_sparse_mqa_logits.hpp" #include "../jit_kernels/impls/sm90_fp8_mqa_logits.hpp" #include "../jit_kernels/impls/sm120_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, @@ -29,7 +32,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); @@ -86,18 +89,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)); @@ -113,8 +116,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)); @@ -124,7 +126,7 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple(weights); DG_HOST_ASSERT(seq_len == _seq_len and num_heads == _num_heads); @@ -153,10 +155,10 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tupleget_arch_major() == 12) ? 128 : 256; - const int block_q = block_qh / num_heads; - DG_HOST_ASSERT(block_qh % num_heads == 0); + DG_HOST_ASSERT(num_heads > 0 and num_heads <= 128 and num_heads % 4 == 0); + // SM120 tiles 128 KV rows per block; SM90/SM100 tile 256 + const int block_kv = (arch_major == 12) ? 128 : 256; + const int block_q = 128 / num_heads; torch::Tensor logits; int aligned_seq_len = align(seq_len, block_q), stride_logits; @@ -181,26 +183,33 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple FP16 MMA accumulator (Q*K score + per-head reduction in FP16); see note above + DG_HOST_ASSERT(not schedule_meta.has_value()); sm100_mqa_logits_f16_weights( q_fp, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits, logits_dtype, seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, block_q, block_kv); + 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); } else if (arch_major == 10) { DG_HOST_ASSERT(qk_dtype == torch::kFloat8_e4m3fn or qk_dtype == kPackedFP4); - DG_HOST_ASSERT(num_heads == 8 or num_heads == 16 or num_heads == 32 or num_heads == 64); sm100_mqa_logits(q_fp, q_sf, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits, logits_dtype, seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, block_q, block_kv, - is_mx_sf, qk_dtype); + clean_logits, is_mx_sf, qk_dtype, schedule_meta); } else if (arch_major == 9) { + DG_HOST_ASSERT(not schedule_meta.has_value()); DG_HOST_ASSERT(not is_mx_sf); DG_HOST_ASSERT(qk_dtype == torch::kFloat8_e4m3fn); DG_HOST_ASSERT(num_heads == 32 or num_heads == 64); DG_HOST_ASSERT(weights.scalar_type() == torch::kFloat); sm90_fp8_mqa_logits(q_fp, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits, logits_dtype, - seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, block_q, block_kv); + seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, block_q, block_kv, + clean_logits); } else if (arch_major == 12) { + DG_HOST_ASSERT(not schedule_meta.has_value()); DG_HOST_ASSERT(qk_dtype == torch::kFloat8_e4m3fn or qk_dtype == kPackedFP4); DG_HOST_ASSERT(num_heads == 16 or num_heads == 32 or num_heads == 64); DG_HOST_ASSERT(weights_dtype == torch::kFloat); @@ -209,13 +218,215 @@ 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. + 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; +} + +// 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; +} + +// 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; - // 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); + 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) % 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); + 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; } @@ -235,7 +446,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(); const int* indices_ptr = nullptr; if (is_varlen) { const auto& indices_tensor = indices.value(); @@ -288,17 +499,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)); @@ -315,8 +526,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 == 32 or block_kv == 64)) or @@ -343,7 +553,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); @@ -397,7 +607,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 ff3ac590c0..08afca93c0 100644 --- a/csrc/apis/einsum.hpp +++ b/csrc/apis/einsum.hpp @@ -1,15 +1,16 @@ #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/sm120_bmk_bnk_mn.hpp" @@ -18,11 +19,9 @@ #include "../jit_kernels/impls/sm120_bf16_gemm.hpp" #include "../jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.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 @@ -51,7 +50,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 == 12) { @@ -63,7 +62,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); @@ -74,8 +73,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); @@ -88,7 +87,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); @@ -99,8 +98,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); @@ -113,12 +112,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); @@ -130,56 +150,75 @@ 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; // AB-swap for small-M decode: BLOCK_M >= 64 wastes lanes at M <= 32, so swap A<->B to // put the small dim on N (BLOCK_N 16/32). Done before the SF transform; the kernel // writes back to the caller's buffer via runtime stride_cd_m/n (see sm120_fp8_fp4_bmm). // Excluded when accumulating (c): swapped strides break the batched epilogue. - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); constexpr int kSwapAbMMax = 32; const bool swap_ab_eligible = arch_major == 12 and m >= 1 and m <= kSwapAbMMax and major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K - and d.stride(-1) == 1 + and d_tensor.stride(-1) == 1 and not c.has_value(); if (swap_ab_eligible) { @@ -197,7 +236,7 @@ static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, sfb, sfa, /*m=*/n, /*n=*/m, k, swap_recipe, swap_recipe_a, swap_recipe_b, batch_size, batch_size, false); sm120_fp8_fp4_bmm( - b, transformed_sfa_swap, a, transformed_sfb_swap, c, d, + b, transformed_sfa_swap, a, transformed_sfb_swap, c, d_tensor, batch_size, /*m=*/n, /*n=*/m, k, gran_k_a_swap, gran_k_b_swap, major_b, major_a, compiled_dims, @@ -211,51 +250,62 @@ static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, // Dispatch implementation if (arch_major == 12) { - sm120_fp8_fp4_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); + sm120_fp8_fp4_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); } else 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") { // (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}); auto perm_b = b.first.permute({0, 2, 1}); - auto perm_sfb = b.second.permute({0, 2, 1}); + const auto perm_sfb = b.second.permute({0, 2, 1}); // SM120: B is MN-major after permute; .contiguous() to K-major (scalar MN-major path ~3x slower). if (arch_major == 12) { perm_b = perm_b.contiguous(); } - 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") { // (batch_size, m, n, k): (h, d, r, b) + DG_HOST_ASSERT(std::holds_alternative(d)); auto perm_a = a.first.permute({1, 2, 0}); - auto perm_sfa = a.second.permute({1, 2, 0}); + const auto perm_sfa = a.second.permute({1, 2, 0}); auto perm_b = b.first.permute({1, 2, 0}); - auto perm_sfb = b.second.permute({1, 2, 0}); + const auto perm_sfb = b.second.permute({1, 2, 0}); // SM120: A/B MN-major after permute; force K-major (MN-major A unsupported, scalar path ~3x slower). if (arch_major == 12) { perm_a = perm_a.contiguous(); @@ -263,22 +313,17 @@ static void fp8_einsum(const std::string& expr, } 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 902ae699a9..649d559c17 100644 --- a/csrc/apis/gemm.hpp +++ b/csrc/apis/gemm.hpp @@ -2,7 +2,6 @@ #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" @@ -10,7 +9,6 @@ #include "../jit_kernels/impls/sm100_bf16_gemm.hpp" #include "../jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp" #include "../jit_kernels/impls/sm120_bf16_gemm.hpp" -#endif #include "../jit_kernels/impls/smxx_cublaslt.hpp" @@ -19,16 +17,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()); @@ -38,6 +42,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; } @@ -47,7 +54,6 @@ static bool early_return(const int& m, const int &n, const int& k, return false; } -#if DG_TENSORMAP_COMPATIBLE // SM120 MMA consumes K-major operands: repack packed-FP4 (needs `logical_mn`) or copy static torch::Tensor sm120_to_k_major(const torch::Tensor& t, const cute::UMMA::Major& major, const int& logical_mn) { @@ -55,7 +61,6 @@ static torch::Tensor sm120_to_k_major(const torch::Tensor& t, const cute::UMMA:: return t; return t.scalar_type() == kPackedFP4 ? fp4_repack_to_k_major(t, logical_mn) : t.contiguous(); } -#endif static int check_k_grouped_args(const std::optional>& ks_cpu, const torch::Tensor& grouped_layout, @@ -80,9 +85,8 @@ 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 - // SM120: AB-swap decision must precede the single SF transform, so it owns its own flow +// NOTES: the SM120 kernel has no alpha epilogue, callers must assert `alpha` is unset static void fp8_fp4_gemm_nt_sm120(const std::pair& a, const std::pair& b, const torch::Tensor& d, @@ -162,11 +166,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); } @@ -175,7 +180,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); @@ -193,6 +198,7 @@ static void fp8_fp4_gemm_nt(const std::pair& a, a.second, b.second, m, n, k, recipe, recipe_a, recipe_b, std::nullopt, std::nullopt, disable_ue8m0_cast); 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); @@ -202,11 +208,12 @@ 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"); } } else if (arch_major == 12) { + DG_HOST_ASSERT(not alpha.has_value() and "FP8 GEMM alpha requires SM100"); fp8_fp4_gemm_nt_sm120(a, b, d, c, recipe, recipe_a, recipe_b, compiled_dims, disable_ue8m0_cast, major_a, major_b, m, n, k); } else { @@ -222,9 +229,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, @@ -235,10 +243,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, @@ -249,9 +258,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, @@ -270,12 +280,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); @@ -363,7 +373,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); @@ -411,16 +421,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__); @@ -428,28 +442,28 @@ 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 if (arch_major == 12) { DG_HOST_ASSERT(not use_psum_layout and ks_cpu.has_value() and not ks_cpu.value().empty()); + 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); // SM120: single transpose [sum_k, M/N] → [M/N, sum_k] with constant stride=sum_k. // Kernel uses kKGroupedConstantStride: per-group only replaces addr+dim, not stride. const auto a_k = a.first.t().contiguous(); const auto b_k = b.first.t().contiguous(); - const auto num_sms = device_runtime->get_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)); sm120_k_grouped_fp8_fp4_gemm_1d1d(a_k, sfa, b_k, sfb, c, d, m, n, @@ -489,7 +503,6 @@ static void k_grouped_fp8_gemm_nt_contiguous(const std::pair(sum_k) * m); DG_HOST_ASSERT(sum_nk == static_cast(sum_k) * n); - const auto arch_major = device_runtime->get_arch_major(); // Contiguity checks DG_HOST_ASSERT(a.first.is_contiguous()); @@ -507,11 +520,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 = 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); @@ -524,14 +538,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); @@ -553,12 +623,19 @@ 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 if (arch_major == 12) { + DG_HOST_ASSERT(not alpha.has_value() and "BF16 GEMM alpha requires SM100"); sm120_bf16_gemm(sm120_to_k_major(a, major_a, m), sm120_to_k_major(b, major_b, n), c, d, m, n, k, cute::UMMA::Major::K, cute::UMMA::Major::K, compiled_dims); } else { @@ -570,24 +647,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, @@ -631,7 +711,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, @@ -684,7 +764,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); @@ -713,7 +793,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__); @@ -722,16 +802,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); @@ -746,8 +827,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` @@ -782,34 +861,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, @@ -846,35 +1006,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", @@ -895,8 +1066,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); @@ -906,6 +1075,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 a695f5938e..0f57ab0541 100644 --- a/csrc/apis/hyperconnection.hpp +++ b/csrc/apis/hyperconnection.hpp @@ -2,15 +2,12 @@ #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" #include "../jit_kernels/impls/sm120_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, @@ -48,7 +45,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 == 12) { sm120_tf32_hc_prenorm_gemm(a, b, d, sqr_sum, m, n, k, num_splits.has_value() ? num_splits.value() : 1); } else if (arch_major == 9) { @@ -60,14 +57,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 5f5c870a62..0c5460de4b 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) @@ -128,17 +124,14 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te sf_input, grouped_layout, ks_cpu, gran_k, k_alignment, use_psum_layout); } - // INT (already packed UE8M0) on SM100/SM120 - if (sf.scalar_type() == torch::kInt and (arch_major == 10 or arch_major == 12)) + // Pre-packed UE8M0: only gran_k=32 on SM100; any supported granularity on SM120 + if (sf.scalar_type() == torch::kInt and ((arch_major == 10 and gran_k == 32) or arch_major == 12)) 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, @@ -153,8 +146,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..c02e57cf72 --- /dev/null +++ b/csrc/apis/mega_gate.hpp @@ -0,0 +1,238 @@ +#pragma once + +#include +#include + +#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 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); + 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); + // 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; + 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 93a9138ce4..40a09c8f74 100644 --- a/csrc/apis/mega.hpp +++ b/csrc/apis/mega_moe.hpp @@ -8,10 +8,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" #include "../jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp" @@ -48,7 +45,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; @@ -221,13 +218,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_); @@ -279,7 +277,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_); @@ -315,7 +314,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_(); } @@ -364,7 +363,7 @@ static void fp4_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(); DG_HOST_ASSERT(l1_weights.scalar_type() == kPackedFP4 and l2_weights.scalar_type() == kPackedFP4); const auto [num_experts_per_rank, intermediate_hidden_2, hidden] = check_grouped_ab_fp8_fp4(l1_weights, cute::UMMA::Major::K, arch_major); @@ -510,7 +509,7 @@ static void fp4_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_(); } @@ -542,7 +541,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); @@ -618,19 +617,17 @@ 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("fp4_fp4_mega_moe", &fp4_fp4_mega_moe); m.def("bf16_mega_moe", &bf16_mega_moe); -#endif } } // namespace deep_gemm::mega diff --git a/csrc/apis/sm90_mega.hpp b/csrc/apis/sm90_mega.hpp index a7cee9469d..983c153dc3 100644 --- a/csrc/apis/sm90_mega.hpp +++ b/csrc/apis/sm90_mega.hpp @@ -8,13 +8,11 @@ #include #include -#if DG_TENSORMAP_COMPATIBLE -#include "../jit/compiler.hpp" -#endif -#include "../jit/device_runtime.hpp" +#include + +#include "../runtime/runtime.hpp" #include "../jit_kernels/impls/sm90_fp8_mega_moe.hpp" #include "../utils/layout.hpp" -#include "../utils/system.hpp" namespace deep_gemm::mega { @@ -168,7 +166,7 @@ static void fp8_mega_moe( const auto [l2_weights, l2_weights_sf] = l2_weights_tuple; // Architecture check - const auto arch_major = device_runtime->get_arch_major(); + const auto arch_major = jit->device.get_arch_major(); if (arch_major != 9) DG_HOST_UNREACHABLE("SM90 FP8 MegaMoE requires a compute capability 9.x GPU"); @@ -252,16 +250,14 @@ static void fp8_mega_moe( hidden, intermediate_hidden, activation_clamp, fast_math); - if (get_env("DG_COMM_KERNEL_DEBUG")) + if (deep_jit::get_env("DG_COMM_KERNEL_DEBUG")) sym_buffer.zero_(); } static void register_sm90_apis(pybind11::module_& m) { -#if DG_TENSORMAP_COMPATIBLE m.def("get_token_alignment_for_sm90_mega_moe", &get_token_alignment_for_sm90_mega_moe); m.def("get_symm_buffer_size_for_sm90_mega_moe", &get_symm_buffer_size_for_sm90_mega_moe); m.def("fp8_mega_moe", &fp8_mega_moe); -#endif } } // namespace deep_gemm::mega diff --git a/csrc/indexing/main.cu b/csrc/indexing/main.cu index 6736207aa3..395de0f980 100644 --- a/csrc/indexing/main.cu +++ b/csrc/indexing/main.cu @@ -26,6 +26,7 @@ // 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 42beaf51ca..0000000000 --- a/csrc/jit/compiler.hpp +++ /dev/null @@ -1,369 +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 - device_runtime->set_support_arch_family(nvcc_major > 12 or nvcc_minor >= 9); - const auto arch = device_runtime->get_arch(false); - // SM120a requires -gencode (--gpu-architecture makes ptxas fall back to sm_120, - // losing block_scale and other arch-specific features) - const auto arch_flag = device_runtime->get_arch_major() == 12 - ? fmt::format("-gencode=arch=compute_{},code=sm_{}", arch, arch) - : fmt::format("--gpu-architecture=sm_{}", arch); - flags = fmt::format("{} -I{} {} " - "--compiler-options=-fPIC,-O3,-fconcepts,-Wno-deprecated-declarations,-Wno-abi " - "-O3 --expt-relaxed-constexpr --expt-extended-lambda", - flags, library_include_path.c_str(), arch_flag); - } - - 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 - device_runtime->set_support_arch_family(major > 12 or minor >= 9); - const auto arch = device_runtime->get_arch(false); - 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 d433558bfc..0000000000 --- a/csrc/jit/device_runtime.hpp +++ /dev/null @@ -1,153 +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; - bool support_arch_family = 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}; - } - - // Set by the active JIT compiler at construction (NVCC / NVRTC >= 12.9). - void set_support_arch_family(const bool& value) { - support_arch_family = value; - } - - bool get_support_arch_family() const { - return support_arch_family; - } - - std::string get_arch(const bool& number_only = 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"; - } - // SM12.1 reuses the SM120 family cubin when `sm_f` is available. - if (major == 12 and (minor != 1 or support_arch_family)) { - if (number_only) - return "120"; - return support_arch_family ? "120f" : "120a"; - } - 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 7b7c94b936..ad9f6a35b3 100644 --- a/csrc/jit_kernels/heuristics/common.hpp +++ b/csrc/jit_kernels/heuristics/common.hpp @@ -1,22 +1,19 @@ #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 { -inline int get_byte_addressable_element_size(const MmaKind& mma_kind) { - const int element_size = get_element_size(mma_kind); - DG_HOST_ASSERT(element_size != -1 and "Unknown MMA kind"); - DG_HOST_ASSERT(element_size > 0); - return element_size; -} - template static GemmConfig get_best_config(const GemmDesc& desc) { desc.check_validity(); @@ -44,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 8e245757d8..5c094fa50a 100644 --- a/csrc/jit_kernels/heuristics/config.hpp +++ b/csrc/jit_kernels/heuristics/config.hpp @@ -39,8 +39,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 { @@ -50,7 +60,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); } @@ -117,11 +130,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 af9f6decff..9225223345 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 { @@ -63,14 +67,25 @@ static MmaKind parse_mma_kind(const std::string& mma_type_str) { return MmaKind::BF16; if (mma_type_str == "fp4xfp4") return MmaKind::NVFP4; - DG_HOST_ASSERT(mma_type_str == "fp8xfp4"); + DG_HOST_ASSERT(mma_type_str == "fp8xfp4" or mma_type_str == "fp8xfp8"); return MmaKind::MXFP8FP4; } +// Element sizes are in bits so that packed NVFP4 (4 bits) is representable static int get_num_mma_elem_bits(const MmaKind& mma_kind) { - if (mma_kind == MmaKind::NVFP4) - return 4; - return get_element_size(mma_kind) * 8; + switch (mma_kind) { + case MmaKind::BF16: return 16; + case MmaKind::MXFP8FP4: return 8; + case MmaKind::NVFP4: return 4; + default: DG_HOST_UNREACHABLE("Unsupported MegaMoE MMA kind"); + } +} + +// Byte-addressable element size; NVFP4 is sub-byte and must go through the bit-based helpers +static int get_num_mma_elem_bytes(const MmaKind& mma_kind) { + const int num_bits = get_num_mma_elem_bits(mma_kind); + DG_HOST_ASSERT(num_bits % 8 == 0); + return num_bits / 8; } static bool is_mma_with_sf(const MmaKind& mma_kind) { @@ -87,32 +102,29 @@ 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` above is in bytes: convert to elements. NVFP4 currently reuses - // the FP8 x FP4 token thresholds without NVFP4-specific threshold tuning; - // it keeps the byte footprint while doubling K elements relative to FP8. - block_k = block_k * 8 / get_num_mma_elem_bits(mma_kind); + } + const int store_block_m = block_m <= 16 ? 8 : block_m <= 64 ? 16 : block_m <= 192 ? 32 : 40; + // `block_k` is 128 bytes worth of elements. NVFP4 currently reuses the FP8 x FP4 + // token thresholds without NVFP4-specific tuning; it keeps the byte footprint + // while doubling K elements relative to FP8. + const int block_k = 128 * 8 / get_num_mma_elem_bits(mma_kind); // Check whether our `block_m` lies in `kCandidateBlockM` DG_HOST_ASSERT(std::any_of( @@ -120,8 +132,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( @@ -205,12 +217,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, mma_kind) : std::pair(0, 0); - // NOTES: FP8 activations and FP4 weights (unpacked to 8-bit in smem, or packed for NVFP4) + // NOTES: FP8 activations, FP8 weights, and FP4 weights (unpacked to 8-bit in smem, or packed for NVFP4) // all use 128B swizzle const int swizzle_acts_mode = 128; const int swizzle_weights_mode = 128; @@ -220,10 +233,11 @@ 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 + // NOTES: NVFP4 keeps the 4 KiB pull granularity its kernel was tuned with + const int num_max_pull_bytes = mma_kind == MmaKind::MXFP8FP4 ? 8192 : 4096; int num_bytes_per_pull = hidden * get_num_mma_elem_bits(mma_kind) / 8; - 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; } @@ -249,8 +263,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 a06b7b7bbe..3493f13676 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; } @@ -44,36 +54,33 @@ class HeuristicsRuntime { return mk_alignment_for_contiguous_layout; } - // Per-arch BLOCK_M search: start at `max_block_m`, step down by `step` (never below - // `min_block_m`) until the tile no longer over-covers M. - struct ContiguousMKAlignment { int max_block_m, min_block_m, step; }; - - static ContiguousMKAlignment get_contiguous_mk_alignment(const int& arch_major) { - // SM120: warp layout is kMWarps(4) * MMA_M(16), so BLOCK_M is a multiple of 64 - if (arch_major == 12) - return {128, 64, 64}; - // SM100: follow the upstream 32-row MMA steps from 224 down to 32 - if (arch_major == 10) - return {224, 32, 32}; - // SM90 and others: fixed legacy alignment, no shrinking - return {kLegacyMKAlignmentForContiguousLayout, kLegacyMKAlignmentForContiguousLayout, 1}; - } - static int get_theoretical_mk_alignment_for_contiguous_layout(const std::optional& expected_m, const std::optional& num_groups = std::nullopt) { - const auto spec = get_contiguous_mk_alignment(device_runtime->get_arch_major()); - int block_m = spec.max_block_m; - if (expected_m.has_value()) { - // Grouped layouts must cover the per-group M, not the summed M - int per_group_m = expected_m.value(); - if (num_groups.has_value() and num_groups.value() > 0) - per_group_m = (per_group_m + num_groups.value() - 1) / num_groups.value(); - for (; block_m > spec.min_block_m and block_m - spec.step >= per_group_m; block_m -= spec.step); + const auto arch_major = jit->device.get_arch_major(); + if (arch_major == 10) { + // 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; + } + if (arch_major == 12) { + // SM120: warp layout is kMWarps(4) * MMA_M(16), so BLOCK_M is a multiple of 64. + // Start at 128 and step down to 64 while the tile still over-covers the per-group M. + constexpr int kMaxBlockM = 128, kMinBlockM = 64, kStep = 64; + int block_m = kMaxBlockM; + if (expected_m.has_value()) { + // Grouped layouts must cover the per-group M, not the summed M + int per_group_m = expected_m.value(); + if (num_groups.has_value() and num_groups.value() > 0) + per_group_m = (per_group_m + num_groups.value() - 1) / num_groups.value(); + for (; block_m > kMinBlockM and block_m - kStep >= per_group_m; block_m -= kStep); + } + return block_m; } - return block_m; + // SM90 and others: fixed legacy alignment + return kLegacyMKAlignmentForContiguousLayout; } }; -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 068ed892c3..f2b0e9e182 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,16 +19,18 @@ 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::BF16: return {0, 0}; case MmaKind::MXFP8FP4: - case MmaKind::NVFP4: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)}; + case MmaKind::MXF4: + case MmaKind::NVFP4: + 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_byte_addressable_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 @@ -123,9 +126,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}; @@ -144,6 +148,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; } @@ -159,11 +174,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)); @@ -174,26 +191,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; @@ -201,8 +235,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 @@ -213,7 +247,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/sm120.hpp b/csrc/jit_kernels/heuristics/sm120.hpp index 4e7b72d52f..0c6789d9df 100644 --- a/csrc/jit_kernels/heuristics/sm120.hpp +++ b/csrc/jit_kernels/heuristics/sm120.hpp @@ -1,5 +1,10 @@ #pragma once +#include +#include +#include +#include + #include #include @@ -13,10 +18,16 @@ namespace deep_gemm { struct SM120ArchSpec { static constexpr int smem_capacity = 101376; // 99KB + // SM120 heuristics address operands in whole bytes: packed FP4 (the `MXF4` MMA kind, 4 bits) + // is tiled like a 1-byte element, with `BLOCK_K` counted in logical elements + static int get_byte_addressable_element_size(const GemmDesc& desc) { + return std::max(get_num_element_bits(desc.get_mma_kind()) / 8, 1); + } + static constexpr int kMinBlockM = 64; // kMWarps(4) * MMA_M(16), both FP8 and BF16 with kNWarps=2 static std::vector get_layout_candidates(const GemmDesc& desc) { - const int elem_size = get_byte_addressable_element_size(desc.get_mma_kind()); + const int elem_size = get_byte_addressable_element_size(desc); const int runtime_align = heuristics_runtime->get_mk_alignment_for_contiguous_layout(); const int expected_m = desc.get_expected_m(); @@ -242,7 +253,7 @@ struct SM120ArchSpec { const int64_t expected_k = desc.get_expected_k(); const int k_blocks = ceil_div(static_cast(expected_k), layout.block_k); - const int elem_size = get_byte_addressable_element_size(desc.get_mma_kind()); + const int elem_size = get_byte_addressable_element_size(desc); const int sf_bytes_a = (desc.kernel_type == KernelType::Kernel1D1D) ? align(layout.block_m * 4, 128) : 0; const int sf_bytes_b = (desc.kernel_type == KernelType::Kernel1D1D) diff --git a/csrc/jit_kernels/heuristics/sm90.hpp b/csrc/jit_kernels/heuristics/sm90.hpp index 4735aac437..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_byte_addressable_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/sm90_mega_moe.hpp b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp index b85b8ed599..da163c10a6 100644 --- a/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp +++ b/csrc/jit_kernels/heuristics/sm90_mega_moe.hpp @@ -3,6 +3,8 @@ #include #include #include +#include +#include #include #include @@ -11,9 +13,9 @@ #include #include #include +#include #include "../../utils/math.hpp" -#include "../../utils/system.hpp" #include "sm90.hpp" namespace deep_gemm { @@ -889,8 +891,8 @@ static Sm90MoeLaunchConfig select_mega_moe_sm90( is_sm90_moe_launch_config_legal(input, candidate, tuning.selected)) { result = candidate; } - if (get_env("DG_JIT_DEBUG") or get_env("DG_PRINT_CONFIGS")) { - const auto key = fmt::format( + if (deep_jit::get_env("DG_JIT_DEBUG") or deep_jit::get_env("DG_PRINT_CONFIGS")) { + const auto key = std::format( "Sm90MoeLaunchConfig(num_ranks={}, num_experts={}, hidden={}, intermediate_hidden={}, num_max_tokens_per_rank={}, num_tokens={}, num_topk={})", input.num_ranks, input.num_experts, input.hidden, input.intermediate_hidden, input.num_max_tokens_per_rank, diff --git a/csrc/jit_kernels/heuristics/utils.hpp b/csrc/jit_kernels/heuristics/utils.hpp index 17d2ae07e6..b5a1aa9495 100644 --- a/csrc/jit_kernels/heuristics/utils.hpp +++ b/csrc/jit_kernels/heuristics/utils.hpp @@ -9,6 +9,16 @@ 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; + case MmaKind::NVFP4: 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 b2959494a5..94e3d98f95 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()"; @@ -86,21 +95,17 @@ static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType& case torch::kFloat16: return CU_TENSOR_MAP_DATA_TYPE_FLOAT16; // Raw bytes (e.g. packed-FP4 tensors moved without smem unpacking) case torch::kUInt8: 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) { @@ -121,6 +126,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; @@ -138,13 +146,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), @@ -160,6 +168,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; @@ -177,12 +190,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), 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_fp4_fp4_mega_moe.hpp b/csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp index d0db58518c..4c86193526 100644 --- a/csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp4_fp4_mega_moe.hpp @@ -1,131 +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 SM100FP4FP4MegaMoERuntime 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; - float routed_scaling_factor; - const void* l1_alphas; - const void* l2_alphas; - const void* a2_scales; - 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_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_fp4_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.routed_scaling_factor, - args.l1_alphas, args.l2_alphas, args.a2_scales, - 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_weights, - args.tensor_map_shared_l1_output, - args.tensor_map_shared_l2_acts, - args.tensor_map_shared_l2_weights - )); - } -}; - static void sm100_fp4_fp4_mega_moe( const torch::Tensor& y, const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf, @@ -263,44 +150,77 @@ static void sm100_fp4_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 SM100FP4FP4MegaMoERuntime::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, - .routed_scaling_factor = routed_scaling_factor, - .l1_alphas = l1_alphas, .l2_alphas = l2_alphas, .a2_scales = a2_scales, - .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_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_fp4_fp4_mega_moe", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_fp4_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")); - const auto code = SM100FP4FP4MegaMoERuntime::generate(args); - const auto runtime = compiler->build("sm100_fp4_fp4_mega_moe", code); - SM100FP4FP4MegaMoERuntime::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, + routed_scaling_factor, + l1_alphas, l2_alphas, a2_scales, + 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_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_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 b323baa3c1..a92c2f28c1 100644 --- a/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp @@ -1,140 +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; - bool use_situ; - float situ_beta, situ_linear_beta; - 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", - args.use_situ ? "true" : "false", - to_string(args.situ_beta), to_string(args.situ_linear_beta)); - } - - 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, @@ -193,7 +71,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). @@ -283,50 +161,88 @@ 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, - .use_situ = use_situ, - .situ_beta = situ_beta, - .situ_linear_beta = situ_linear_beta, - .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 + // NOTES: the template order follows the device kernel: `kFastMath`, then the SiTU + // parameters (`kUseSiTU`, `kSiTUBeta`, `kSiTULinearBeta`), then `weight_dtype_t` + const auto kernel = jit->compile("sm100_fp8_fp4_mega_moe", std::format(R"( +#include + +using namespace deep_gemm; + +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", + use_situ ? "true" : "false", + to_string(situ_beta), to_string(situ_linear_beta), + to_string(l1_weights.scalar_type()))); - const auto code = SM100FP8FP4MegaMoERuntime::generate(args); - const auto runtime = compiler->build("sm100_fp8_fp4_mega_moe", code); - SM100FP8FP4MegaMoERuntime::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_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_mqa_logits_f16_weights.hpp b/csrc/jit_kernels/impls/sm100_mqa_logits_f16_weights.hpp index 8e39e7cba0..1d031c3214 100644 --- a/csrc/jit_kernels/impls/sm100_mqa_logits_f16_weights.hpp +++ b/csrc/jit_kernels/impls/sm100_mqa_logits_f16_weights.hpp @@ -1,81 +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/sm100.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM100MQALogitsF16WeightsRuntime 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; - - uint32_t* cu_seq_len_k_start_and_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(&sm100_fp8_mqa_logits_f16_weights< - {}, {}, - {}, - {}, {}, - {}, {}, - {}, {}, - {} - >); -}}; -)", 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.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, - static_cast(args.stride_logits), - args.cu_seq_len_k_start_and_end, args.logits, args.tensor_map_q, - args.tensor_map_kv, args.tensor_map_kv_scales, - args.tensor_map_weights)); - } -}; - static void sm100_mqa_logits_f16_weights( const torch::Tensor& q, const torch::Tensor& kv, const torch::Tensor& kv_scales, const torch::Tensor& weights, @@ -85,7 +17,8 @@ static void sm100_mqa_logits_f16_weights( 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) { - DG_HOST_ASSERT(device_runtime->get_arch_major() == 10); + DG_HOST_ASSERT(jit->device.get_arch_major() == 10); + DG_HOST_ASSERT(128 % num_heads == 0); constexpr int num_specialized_threads = 128; constexpr int num_q_stages = 5, num_kv_stages = 8; @@ -144,36 +77,43 @@ static void sm100_mqa_logits_f16_weights( valid_offsets.select(1, 1).copy_(cu_seq_len_k_end); cu_seq_len_k_start_and_end = cu_seq_len_k_start_and_end.reshape({-1}).contiguous(); - const SM100MQALogitsF16WeightsRuntime::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_and_end = reinterpret_cast( - cu_seq_len_k_start_and_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, 2) - }; - const auto code = SM100MQALogitsF16WeightsRuntime::generate(args); - const auto runtime = - compiler->build("sm100_mqa_logits_f16_weights", code); - SM100MQALogitsF16WeightsRuntime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("sm100_mqa_logits_f16_weights", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm100_fp8_mqa_logits_f16_weights< + {}, {}, + {}, + {}, {}, + {}, {}, + {}, {}, + {} + >); +}}; +)", num_heads, head_dim, + is_compressed_logits, + block_q, block_kv, + num_q_stages, num_kv_stages, + num_specialized_threads, num_math_threads, + to_string(logits_dtype))); + + // Launch (two-CTA cluster) + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(runtime->get_num_sms(), 1, 1), + .block_dim = dim3(num_specialized_threads + num_math_threads, 1, 1), + .cluster_dim = dim3(2, 1, 1), + }, + seq_len, seq_len_kv, max_seqlen_k, + static_cast(stride_logits), + reinterpret_cast(cu_seq_len_k_start_and_end.data_ptr()), + logits.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/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/sm120_bf16_gemm.hpp b/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp index e272caf346..d672828e4c 100644 --- a/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_bf16_gemm.hpp @@ -1,26 +1,29 @@ #pragma once +#include +#include +#include +#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/sm120.hpp" +#include "epilogue.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM120BF16GemmRuntime final: public LaunchRuntime { +class SM120BF16GemmRuntime final { public: struct Args { GemmDesc gemm_desc; GemmConfig gemm_config; - LaunchArgs launch_args; - const std::optional epilogue_type; + deep_jit::cuda::LaunchOptions options; + std::optional epilogue_type; void* gmem_d; void* gmem_c; @@ -33,12 +36,12 @@ class SM120BF16GemmRuntime final: public LaunchRuntime { CUtensorMap tensor_map_cd; }; - static std::string generate_impl(const Args& args) { + static void compile_and_launch(const std::string& tag, const Args& args) { // kNWarps: BM=64 needs kNWarps=2 (kMWarps=4); BM=128 uses default // kNWarps=kNumMathWarps (all warps along M, original layout). const uint32_t block_m = args.gemm_config.layout.block_m; const uint32_t k_n_warps = (block_m % 64 == 0 and block_m < 128) ? 2u : 1u; - return fmt::format(R"( + const auto kernel = jit->compile(tag, std::format(R"( #include using namespace deep_gemm; @@ -75,18 +78,19 @@ static void __instantiate_kernel() {{ to_string(args.gemm_desc.cd_dtype), get_default_epilogue_type(args.epilogue_type), args.gemm_desc.major_b == cute::UMMA::Major::K ? "true" : "false", - k_n_warps); - } + k_n_warps)); - 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_d, args.gmem_c, 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, args.tensor_map_b, - args.tensor_map_cd)); + args.tensor_map_cd + ); } }; @@ -107,8 +111,8 @@ static void sm120_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(), + .num_sms = runtime->get_num_sms(), + .tc_util = runtime->get_tc_util(), .compiled_dims = compiled_dims }; const auto config = get_best_config(desc); @@ -129,12 +133,15 @@ static void sm120_bf16_gemm(const torch::Tensor& a, n, 1, config.storage_config.swizzle_cd_mode); - const SM120BF16GemmRuntime::Args args = { + // Compile and launch + SM120BF16GemmRuntime::compile_and_launch("sm120_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, - 1), + .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), + }, .epilogue_type = std::nullopt, .gmem_d = d.data_ptr(), .gmem_c = c.has_value() ? cd.data_ptr() : nullptr, @@ -145,10 +152,7 @@ static void sm120_bf16_gemm(const torch::Tensor& a, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM120BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm120_bf16_gemm", code); - SM120BF16GemmRuntime::launch(runtime, args); + }); } static void sm120_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, @@ -176,8 +180,8 @@ static void sm120_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(), + .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, @@ -200,12 +204,15 @@ static void sm120_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); - const SM120BF16GemmRuntime::Args args = { + // Compile and launch + SM120BF16GemmRuntime::compile_and_launch("sm120_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, - 1), + .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), + }, .epilogue_type = std::nullopt, .gmem_d = d.data_ptr(), .gmem_c = nullptr, @@ -216,10 +223,7 @@ static void sm120_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM120BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm120_m_grouped_bf16_gemm_contiguous", code); - SM120BF16GemmRuntime::launch(runtime, args); + }); } static void sm120_m_grouped_bf16_gemm_masked(const torch::Tensor& a, @@ -240,8 +244,8 @@ static void sm120_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(), + .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 }; @@ -262,12 +266,15 @@ static void sm120_m_grouped_bf16_gemm_masked(const torch::Tensor& a, static_cast(d.stride(-2)), num_groups, config.storage_config.swizzle_cd_mode); - const SM120BF16GemmRuntime::Args args = { + // Compile and launch + SM120BF16GemmRuntime::compile_and_launch("sm120_m_grouped_bf16_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, - 1), + .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), + }, .epilogue_type = std::nullopt, .gmem_d = d.data_ptr(), .gmem_c = nullptr, @@ -278,10 +285,7 @@ static void sm120_m_grouped_bf16_gemm_masked(const torch::Tensor& a, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM120BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm120_m_grouped_bf16_gemm_masked", code); - SM120BF16GemmRuntime::launch(runtime, args); + }); } static void sm120_bf16_k_grouped_gemm(const torch::Tensor& a, @@ -326,14 +330,14 @@ static void sm120_bf16_k_grouped_gemm(const torch::Tensor& a, .cd_dtype = d.scalar_type(), .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K, .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); // Allocate tensor map buffer for dynamic replacement (A + B per SM) - const auto num_sms = device_runtime->get_num_sms(); + const auto num_sms = runtime->get_num_sms(); const auto tensor_map_buffer = torch::empty( {num_sms * 2 * static_cast(sizeof(CUtensorMap))}, a.options().dtype(torch::kByte)); @@ -359,11 +363,16 @@ static void sm120_bf16_k_grouped_gemm(const torch::Tensor& a, n, num_groups, config.storage_config.swizzle_cd_mode); const auto cd = c.value_or(d); - const SM120BF16GemmRuntime::Args args = { + + // Compile and launch + SM120BF16GemmRuntime::compile_and_launch("sm120_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, 1), + .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), + }, .epilogue_type = std::nullopt, .gmem_d = d.data_ptr(), .gmem_c = cd.data_ptr(), @@ -374,10 +383,7 @@ static void sm120_bf16_k_grouped_gemm(const torch::Tensor& a, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM120BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm120_bf16_k_grouped_gemm", code); - SM120BF16GemmRuntime::launch(runtime, args); + }); } static void sm120_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, @@ -393,8 +399,8 @@ static void sm120_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); @@ -411,12 +417,15 @@ static void sm120_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, tensor_d.stride(0), tensor_d.stride(1), config.storage_config.swizzle_cd_mode); - const SM120BF16GemmRuntime::Args args = { + // Compile and launch + SM120BF16GemmRuntime::compile_and_launch("sm120_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, - 1), + .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), + }, .epilogue_type = std::nullopt, .gmem_d = tensor_d.data_ptr(), .gmem_c = nullptr, @@ -427,10 +436,7 @@ static void sm120_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM120BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm120_bf16_bhr_hdr_bhd", code); - SM120BF16GemmRuntime::launch(runtime, args); + }); } static void sm120_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, @@ -448,8 +454,8 @@ static void sm120_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); @@ -469,12 +475,15 @@ static void sm120_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, tensor_d.stride(0), tensor_d.stride(1), config.storage_config.swizzle_cd_mode); - const SM120BF16GemmRuntime::Args args = { + // Compile and launch + SM120BF16GemmRuntime::compile_and_launch("sm120_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, - 1), + .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), + }, .epilogue_type = std::nullopt, .gmem_d = tensor_d.data_ptr(), .gmem_c = nullptr, @@ -485,10 +494,7 @@ static void sm120_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a, .tensor_map_a = tensor_map_a, .tensor_map_b = tensor_map_b, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM120BF16GemmRuntime::generate(args); - const auto runtime = compiler->build("sm120_bf16_bhd_hdr_bhr", code); - SM120BF16GemmRuntime::launch(runtime, args); + }); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp b/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp index 55365324c3..61c94c9f9a 100644 --- a/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp +++ b/csrc/jit_kernels/impls/sm120_bmk_bnk_mn.hpp @@ -1,67 +1,19 @@ #pragma once +#include +#include +#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/sm120.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM120BmkBnkMnRuntime 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; - 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(&sm120_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.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 sm120_bmn_bnk_mn_gemm(const torch::Tensor &a, const torch::Tensor &b, const torch::Tensor &d, @@ -78,7 +30,7 @@ static void sm120_bmn_bnk_mn_gemm(const torch::Tensor &a, const int swizzle_ab_mode = get_swizzle_mode(block_k, static_cast(a.element_size())); DG_HOST_ASSERT(swizzle_ab_mode == 128); - 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 +52,8 @@ static void sm120_bmn_bnk_mn_gemm(const torch::Tensor &a, } DG_HOST_ASSERT(num_stages > 0); - if (get_env("DG_JIT_DEBUG", 0)) { + // Print configs + if (deep_jit::get_env("DG_PRINT_CONFIGS")) { printf("SM120 bmk_bnk_mn: S: %d, M: %d, N: %d, K: %d -> " "split_factor: %d, stages: %d, shared memory: %d\n", s, m, n, k, split_factor, num_stages, smem_size); @@ -109,22 +62,39 @@ static void sm120_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 SM120BmkBnkMnRuntime::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, - .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 = SM120BmkBnkMnRuntime::generate(args); - const auto runtime = compiler->build("sm120_bmn_bnk_mn_gemm", code); - SM120BmkBnkMnRuntime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("sm120_bmn_bnk_mn_gemm", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_bmn_bnk_mn_gemm_impl< + {}, {}, {}, + {}, {}, {}, + {}, + {}, + {}, + {}, {} + >); +}}; +)", + m, n, k, + block_m, block_n, block_k, + split_factor, + swizzle_ab_mode, + 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/sm120_fp8_fp4_gemm_1d1d.hpp b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp index c8fab5839e..8b398bb13b 100644 --- a/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp +++ b/csrc/jit_kernels/impls/sm120_fp8_fp4_gemm_1d1d.hpp @@ -1,12 +1,14 @@ #pragma once +#include +#include +#include +#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/sm120.hpp" @@ -15,13 +17,13 @@ namespace deep_gemm { -class SM120FP8FP4Gemm1D1DRuntime final: public LaunchRuntime { +class SM120FP8FP4Gemm1D1DRuntime final { public: struct Args { GemmDesc gemm_desc; GemmConfig gemm_config; - LaunchArgs launch_args; - const std::optional epilogue_type; + deep_jit::cuda::LaunchOptions options; + std::optional epilogue_type; int gran_k_a, gran_k_b; bool is_fp4; @@ -46,8 +48,8 @@ class SM120FP8FP4Gemm1D1DRuntime final: public LaunchRuntimecompile(tag, std::format(R"( #include using namespace deep_gemm; @@ -96,11 +98,11 @@ static void __instantiate_kernel() {{ (args.gemm_desc.major_b == cute::UMMA::Major::K) ? "true" : "false", args.k_grouped_constant_stride ? "true" : "false", args.gemm_config.storage_config.store_block_m, - args.gemm_config.split_k_factor); - } + args.gemm_config.split_k_factor)); - 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_d, args.gmem_c, args.gmem_a_ptr, args.gmem_b_ptr, args.grouped_layout, @@ -110,24 +112,20 @@ static void __instantiate_kernel() {{ args.stride_cd_m, args.stride_cd_n, args.stride_cd_batch, args.tensor_map_a, args.tensor_map_b, args.tensor_map_sfa, args.tensor_map_sfb, - args.tensor_map_cd)); + args.tensor_map_cd + ); } }; -class SM120SplitKReduceRuntime final: public LaunchRuntime { -public: - struct Args { - GemmDesc gemm_desc; - GemmConfig gemm_config; - LaunchArgs launch_args; - int stride_cd_m; - int stride_cd_n; - void* gmem_d; - void* workspace; - }; +static void sm120_split_k_reduce(const torch::Tensor& workspace, const torch::Tensor& d, + const int& m, const int& n, const int& split_k, + const int stride_cd_m, const int stride_cd_n) { + const int total = m * n; + const int threads = 256; + const int blocks = ceil_div(total, threads); - static std::string generate_impl(const Args& args) { - return fmt::format(R"( + // Compile + const auto kernel = jit->compile("sm120_split_k_reduce", std::format(R"( #include using namespace deep_gemm; @@ -136,51 +134,19 @@ static void __instantiate_kernel() {{ auto ptr = reinterpret_cast(&sm120_split_k_reduce_impl<{}, {}>); }}; )", - to_string(args.gemm_desc.cd_dtype), - args.gemm_config.split_k_factor); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.gmem_d, args.workspace, - args.gemm_desc.m, args.gemm_desc.n, args.stride_cd_m, args.stride_cd_n)); - } -}; - -static void sm120_split_k_reduce(const torch::Tensor& workspace, const torch::Tensor& d, - const int& m, const int& n, const int& split_k, - const int stride_cd_m, const int stride_cd_n) { - const int total = m * n; - const int threads = 256; - const int blocks = ceil_div(total, threads); - - const auto desc = GemmDesc { - .gemm_type = GemmType::Normal, - .kernel_type = KernelType::KernelNoSF, - .m = m, .n = n, .k = 0, .num_groups = 1, - .a_dtype = torch::kFloat, .b_dtype = torch::kFloat, - .cd_dtype = d.scalar_type(), - .major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K, - .with_accumulation = false, - .num_sms = blocks, - .tc_util = 100, .compiled_dims = "" - }; - const auto config = GemmConfig { - .split_k_factor = split_k - }; - - const SM120SplitKReduceRuntime::Args args = { - .gemm_desc = desc, - .gemm_config = config, - .launch_args = LaunchArgs(blocks, threads, 0, 0), - .stride_cd_m = stride_cd_m, - .stride_cd_n = stride_cd_n, - .gmem_d = d.data_ptr(), - .workspace = workspace.data_ptr(), - }; - const auto code = SM120SplitKReduceRuntime::generate(args); - const auto runtime = compiler->build("sm120_split_k_reduce", code); - SM120SplitKReduceRuntime::launch(runtime, args); + to_string(d.scalar_type()), + split_k)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = 0, + .grid_dim = dim3(blocks, 1, 1), + .block_dim = dim3(threads, 1, 1), + }, + d.data_ptr(), workspace.data_ptr(), + m, n, stride_cd_m, stride_cd_n + ); } static void sm120_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -209,8 +175,8 @@ static void sm120_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, .max_gran_k = std::max(gran_k_a, gran_k_b), .cd_n_contiguous = !swap_ab // strided-store epilogue (AB-swap output is transposed) @@ -251,12 +217,15 @@ static void sm120_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& if (split_k > 1) workspace = torch::empty({split_k, m, n}, d.options().dtype(torch::kFloat)); - const SM120FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM120FP8FP4Gemm1D1DRuntime::compile_and_launch("sm120_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, - 1), + .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), + }, .epilogue_type = epilogue_type, .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, @@ -279,11 +248,7 @@ static void sm120_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 = SM120FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm120_fp8_fp4_gemm_1d1d", code); - SM120FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); if (split_k > 1) { const int reduce_stride_m = swap_ab ? static_cast(d.stride(-1)) : d_stride; @@ -333,8 +298,8 @@ static void sm120_k_grouped_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torc .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, .max_gran_k = std::max(gran_k_a, gran_k_b), .expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups @@ -366,12 +331,15 @@ static void sm120_k_grouped_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torc n, num_groups, config.storage_config.swizzle_cd_mode); - const SM120FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM120FP8FP4Gemm1D1DRuntime::compile_and_launch("sm120_k_grouped_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, - 1), + .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), + }, .epilogue_type = std::nullopt, .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, @@ -394,10 +362,7 @@ static void sm120_k_grouped_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torc .tensor_map_sfa = tensor_map_sfa, .tensor_map_sfb = tensor_map_sfb, .tensor_map_cd = tensor_map_cd, - }; - const auto code = SM120FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm120_k_grouped_fp8_fp4_gemm_1d1d", code); - SM120FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); } static void sm120_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -432,8 +397,8 @@ static void sm120_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, .max_gran_k = std::max(gran_k_a, gran_k_b), .expected_m = expected_m_for_psum_layout.value_or(m), @@ -465,12 +430,15 @@ static void sm120_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, static_cast(d.stride(-2)), 1, config.storage_config.swizzle_cd_mode); - const SM120FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM120FP8FP4Gemm1D1DRuntime::compile_and_launch("sm120_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, - 1), + .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), + }, .epilogue_type = std::nullopt, .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, @@ -493,10 +461,7 @@ static void sm120_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 = SM120FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm120_m_grouped_fp8_fp4_gemm_contiguous_1d1d", code); - SM120FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); } static void sm120_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, const torch::Tensor& sfa, @@ -524,8 +489,8 @@ static void sm120_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, .max_gran_k = std::max(gran_k_a, gran_k_b), .expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups @@ -555,12 +520,15 @@ static void sm120_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, con static_cast(d.stride(-2)), num_groups, config.storage_config.swizzle_cd_mode); - const SM120FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM120FP8FP4Gemm1D1DRuntime::compile_and_launch("sm120_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, - 1), + .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), + }, .epilogue_type = std::nullopt, .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, @@ -583,10 +551,7 @@ static void sm120_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 = SM120FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm120_m_grouped_fp8_fp4_gemm_masked_1d1d", code); - SM120FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); } static void sm120_fp8_fp4_bmm(const torch::Tensor& a, const torch::Tensor& sfa, @@ -615,8 +580,8 @@ static void sm120_fp8_fp4_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, .max_gran_k = std::max(gran_k_a, gran_k_b), .cd_n_contiguous = !swap_ab // strided-store epilogue (AB-swap output is transposed) @@ -645,12 +610,15 @@ static void sm120_fp8_fp4_bmm(const torch::Tensor& a, const torch::Tensor& sfa, d.stride(1), d.stride(0), config.storage_config.swizzle_cd_mode); - const SM120FP8FP4Gemm1D1DRuntime::Args args = { + // Compile and launch + SM120FP8FP4Gemm1D1DRuntime::compile_and_launch("sm120_fp8_fp4_bmm", { .gemm_desc = desc, .gemm_config = config, - .launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads, - config.pipeline_config.smem_size, - 1), + .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), + }, .epilogue_type = std::nullopt, .gran_k_a = gran_k_a, .gran_k_b = gran_k_b, @@ -673,10 +641,7 @@ static void sm120_fp8_fp4_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, - }; - const auto code = SM120FP8FP4Gemm1D1DRuntime::generate(args); - const auto runtime = compiler->build("sm120_fp8_fp4_bmm", code); - SM120FP8FP4Gemm1D1DRuntime::launch(runtime, args); + }); } } // namespace deep_gemm diff --git a/csrc/jit_kernels/impls/sm120_mqa_logits.hpp b/csrc/jit_kernels/impls/sm120_mqa_logits.hpp index 6db1faace7..86b71a1adb 100644 --- a/csrc/jit_kernels/impls/sm120_mqa_logits.hpp +++ b/csrc/jit_kernels/impls/sm120_mqa_logits.hpp @@ -1,65 +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/sm120.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM120FP8MQALogitsRuntime 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, num_kv_stages; - int block_q, 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_tma_threads, 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(&sm120_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_tma_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 sm120_fp8_mqa_logits(const torch::Tensor& q, const torch::Tensor& kv, const torch::Tensor& kv_scales, @@ -75,7 +23,7 @@ static void sm120_fp8_mqa_logits(const torch::Tensor& q, constexpr int num_tma_threads = 128; constexpr int num_q_stages = 2, num_kv_stages = 3; constexpr int num_math_threads = 256; - const int num_sms = device_runtime->get_num_sms(); + const int num_sms = runtime->get_num_sms(); const bool is_compressed_logits = max_seqlen_k > 0; // Split the KV range only when Q blocks underfill SM120. The split blocks @@ -97,7 +45,8 @@ static void sm120_fp8_mqa_logits(const torch::Tensor& q, } } - DG_HOST_ASSERT(device_runtime->get_arch_major() == 12); + DG_HOST_ASSERT(jit->device.get_arch_major() == 12); + DG_HOST_ASSERT(128 % num_heads == 0); 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, @@ -129,79 +78,38 @@ static void sm120_fp8_mqa_logits(const torch::Tensor& q, smem_size += 4; DG_HOST_ASSERT(smem_size <= SM120ArchSpec::smem_capacity); - const SM120FP8MQALogitsRuntime::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_tma_threads = num_tma_threads, - .num_math_threads = num_math_threads, - .launch_args = LaunchArgs({num_sms, kv_splits}, - num_tma_threads + num_math_threads, smem_size) - }; - const auto code = SM120FP8MQALogitsRuntime::generate(args); - const auto runtime = compiler->build("sm120_fp8_mqa_logits", code); - SM120FP8MQALogitsRuntime::launch(runtime, args); -} - -class SM120FP4MQALogitsRuntime final: public LaunchRuntime { -public: - struct Args { - int seq_len, seq_len_kv, max_seqlen_k, stride_logits; - int num_heads, head_dim; - bool is_compressed_logits; - int num_q_stages, num_kv_stages; - int block_q, block_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 logits_dtype; - int num_tma_threads, 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("sm120_fp8_mqa_logits", std::format(R"( +#include using namespace deep_gemm; static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sm120_fp4_mqa_logits< + auto ptr = reinterpret_cast(&sm120_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_tma_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_sf_q, - args.tensor_map_kv, args.tensor_map_sf_kv, args.tensor_map_weights)); - } -}; +)", num_heads, head_dim, is_compressed_logits, + block_q, block_kv, num_q_stages, num_kv_stages, + num_sms, + num_tma_threads, num_math_threads, to_string(logits_dtype))); + + // Launch: grid.y carries the KV splits + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(num_sms, kv_splits, 1), + .block_dim = dim3(num_tma_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 sm120_fp4_mqa_logits(const torch::Tensor& q, const torch::Tensor& sf_q, @@ -220,8 +128,10 @@ static void sm120_fp4_mqa_logits(const torch::Tensor& q, constexpr int num_math_threads = 256; constexpr int num_q_stages = 2, num_kv_stages = 5; const bool is_compressed_logits = max_seqlen_k > 0; + const int num_sms = runtime->get_num_sms(); - DG_HOST_ASSERT(device_runtime->get_arch_major() == 12); + DG_HOST_ASSERT(jit->device.get_arch_major() == 12); + DG_HOST_ASSERT(128 % num_heads == 0); DG_HOST_ASSERT(head_dim == 128); const auto tensor_map_q = make_tma_2d_desc( q, head_dim, seq_len * num_heads, @@ -254,28 +164,37 @@ static void sm120_fp4_mqa_logits(const torch::Tensor& q, smem_barriers; DG_HOST_ASSERT(smem_size <= SM120ArchSpec::smem_capacity); - const SM120FP4MQALogitsRuntime::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_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, - .logits_dtype = logits_dtype, - .num_tma_threads = num_tma_threads, - .num_math_threads = num_math_threads, - .launch_args = LaunchArgs(device_runtime->get_num_sms(), - num_tma_threads + num_math_threads, smem_size) - }; - const auto code = SM120FP4MQALogitsRuntime::generate(args); - const auto runtime = compiler->build("sm120_fp4_mqa_logits", code); - SM120FP4MQALogitsRuntime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("sm120_fp4_mqa_logits", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_fp4_mqa_logits< + {}, {}, {}, + {}, {}, {}, {}, + {}, {}, {}, {} + >); +}}; +)", num_heads, head_dim, is_compressed_logits, + block_q, block_kv, num_q_stages, num_kv_stages, + num_sms, + num_tma_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_tma_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_sf_q, + tensor_map_kv, tensor_map_sf_kv, tensor_map_weights + ); } static void sm120_mqa_logits(const torch::Tensor& q, @@ -309,43 +228,6 @@ static void sm120_mqa_logits(const torch::Tensor& q, } } -class SM120PagedMQALogitsMetadataRuntime final - : public LaunchRuntime { -public: - struct Args { - int aligned_batch_size, split_kv, num_sms; - bool is_varlen; - int batch_size, next_n, num_next_n_atoms; - 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 - -using namespace deep_gemm; - -static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sched::sm120_paged_mqa_logits_metadata< - {}, {}, {}, {} - >); -}}; -)", 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, - static_cast(args.num_next_n_atoms), args.context_lens, - args.indices, args.schedule_metadata)); - } -}; - static void sm120_paged_mqa_logits_metadata( const torch::Tensor& context_lens, const torch::Tensor& schedule_metadata, const int& batch_size, const int& next_n, const int& block_kv, @@ -361,75 +243,34 @@ static void sm120_paged_mqa_logits_metadata( const int smem_size = num_smem_ints * static_cast(sizeof(int)); DG_HOST_ASSERT(smem_size <= SM120ArchSpec::smem_capacity); - const SM120PagedMQALogitsMetadataRuntime::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, - .num_next_n_atoms = num_next_n_atoms, - .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 = SM120PagedMQALogitsMetadataRuntime::generate(args); - const auto runtime = compiler->build("sm120_paged_mqa_logits_metadata", code); - SM120PagedMQALogitsMetadataRuntime::launch(runtime, args); -} - -class SM120FP8PagedMQALogitsRuntime final - : public LaunchRuntime { -public: - struct Args { - int batch_size, next_n, num_heads, head_dim, block_kv; - bool is_context_lens_2d, is_varlen; - int block_table_stride, logits_stride; - int num_q_stages, num_kv_stages, 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_tma_threads, 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("sm120_paged_mqa_logits_metadata", std::format(R"( +#include using namespace deep_gemm; static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sm120_fp8_paged_mqa_logits< - {}, {}, {}, {}, {}, {}, - {}, {}, {}, {}, {}, {} + auto ptr = reinterpret_cast(&sched::sm120_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_tma_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, + static_cast(num_next_n_atoms), + context_lens.data_ptr(), + const_cast(indices_ptr), + schedule_metadata.data_ptr() + ); +} static void sm120_fp8_paged_mqa_logits( const torch::Tensor& q, const torch::Tensor& kv_cache, @@ -447,7 +288,8 @@ static void sm120_fp8_paged_mqa_logits( constexpr int num_q_stages = 2, num_kv_stages = 3; const int num_groups = split_kv / block_kv; const int next_n_atom = (is_varlen or next_n >= 2) ? 2 : 1; - DG_HOST_ASSERT(device_runtime->get_arch_major() == 12); + DG_HOST_ASSERT(jit->device.get_arch_major() == 12); + DG_HOST_ASSERT(128 % num_heads == 0); DG_HOST_ASSERT(block_kv == 64); DG_HOST_ASSERT(split_kv == 128 and logits_stride % split_kv == 0); @@ -489,80 +331,39 @@ static void sm120_fp8_paged_mqa_logits( smem_q_pipe_size + num_groups * smem_kv_pipe_size + 4; DG_HOST_ASSERT(smem_size <= SM120ArchSpec::smem_capacity); - const SM120FP8PagedMQALogitsRuntime::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_tma_threads = num_tma_threads, .num_math_threads = num_math_threads, - .launch_args = LaunchArgs(num_sms, num_tma_threads + num_math_threads, - smem_size) - }; - const auto code = SM120FP8PagedMQALogitsRuntime::generate(args); - const auto runtime = compiler->build("sm120_fp8_paged_mqa_logits", code); - SM120FP8PagedMQALogitsRuntime::launch(runtime, args); -} - -class SM120FP4PagedMQALogitsRuntime final - : public LaunchRuntime { -public: - struct Args { - int batch_size, next_n, num_heads, head_dim, block_kv; - bool is_context_lens_2d, is_varlen; - int block_table_stride, logits_stride; - int num_q_stages, num_kv_stages, split_kv; - 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 logits_dtype; - int num_tma_threads, num_math_threads; - LaunchArgs launch_args; - }; - - static std::string generate_impl(const Args& args) { - return fmt::format(R"( -#include + // Compile + const auto kernel = jit->compile("sm120_fp8_paged_mqa_logits", std::format(R"( +#include using namespace deep_gemm; static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&sm120_fp4_paged_mqa_logits< + auto ptr = reinterpret_cast(&sm120_fp8_paged_mqa_logits< {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {}, {} >); }}; -)", 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_tma_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_sf_q, - args.tensor_map_kv, args.tensor_map_sf_kv, args.tensor_map_weights)); - } -}; +)", 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_tma_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_tma_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 + ); +} static void sm120_fp4_paged_mqa_logits( const torch::Tensor& q, const torch::Tensor& sf_q, @@ -580,7 +381,7 @@ static void sm120_fp4_paged_mqa_logits( constexpr int num_math_threads = 256; constexpr int num_q_stages = 2, num_kv_stages = 3; const int next_n_atom = (is_varlen or next_n >= 2) ? 2 : 1; - DG_HOST_ASSERT(device_runtime->get_arch_major() == 12); + DG_HOST_ASSERT(jit->device.get_arch_major() == 12); DG_HOST_ASSERT(split_kv == 128 and logits_stride % split_kv == 0); DG_HOST_ASSERT(block_kv == 32 or block_kv == 64); DG_HOST_ASSERT(head_dim == 128); @@ -628,28 +429,38 @@ static void sm120_fp4_paged_mqa_logits( smem_q_pipe_size + num_groups * smem_kv_pipe_size + 4; DG_HOST_ASSERT(smem_size <= SM120ArchSpec::smem_capacity); - const SM120FP4PagedMQALogitsRuntime::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_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, - .logits_dtype = logits_dtype, - .num_tma_threads = num_tma_threads, .num_math_threads = num_math_threads, - .launch_args = LaunchArgs(num_sms, num_tma_threads + num_math_threads, - smem_size) - }; - const auto code = SM120FP4PagedMQALogitsRuntime::generate(args); - const auto runtime = compiler->build("sm120_fp4_paged_mqa_logits", code); - SM120FP4PagedMQALogitsRuntime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("sm120_fp4_paged_mqa_logits", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_fp4_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_tma_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_tma_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_sf_q, + tensor_map_kv, tensor_map_sf_kv, tensor_map_weights + ); } static void sm120_paged_mqa_logits( diff --git a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp index 3067baca5a..af6f62cf46 100644 --- a/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp +++ b/csrc/jit_kernels/impls/sm120_tf32_hc_prenorm_gemm.hpp @@ -1,64 +1,16 @@ #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/sm120.hpp" #include "runtime_utils.hpp" namespace deep_gemm { -class SM120TF32HCPrenormGemmRuntime final: public LaunchRuntime { -public: - struct Args { - int m, n, k; - int block_m, block_n, block_k; - int num_splits; - int num_stages; - int num_math_threads, num_tma_threads; - - LaunchArgs launch_args; - - CUtensorMap tensor_map_a; - CUtensorMap tensor_map_b; - float* gmem_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(&sm120_tf32_hc_prenorm_gemm_impl< - {}, {}, - {}, {}, {}, - {}, - {}, - {}, {} - >); -}}; -)", - args.n, args.k, - args.block_m, args.block_n, args.block_k, - args.num_splits, - args.num_stages, - args.num_math_threads, args.num_tma_threads); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.m, args.tensor_map_a, args.tensor_map_b, args.gmem_d, args.sqr_sum)); - } -}; - static void sm120_tf32_hc_prenorm_gemm(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d, @@ -106,21 +58,38 @@ static void sm120_tf32_hc_prenorm_gemm(const torch::Tensor& a, const int num_m_blocks = ceil_div(m, block_m); const int grid_size = num_m_blocks * num_splits; - const SM120TF32HCPrenormGemmRuntime::Args args = { - .m = m, .n = n, .k = k, - .block_m = block_m, .block_n = block_n, .block_k = block_k, - .num_splits = num_splits, - .num_stages = num_stages, - .num_math_threads = num_math_threads, .num_tma_threads = num_tma_threads, - .launch_args = LaunchArgs(grid_size, num_threads, smem_size, 1), - .tensor_map_a = tensor_map_a, - .tensor_map_b = tensor_map_b, - .gmem_d = reinterpret_cast(d.data_ptr()), - .sqr_sum = reinterpret_cast(sqr_sum.data_ptr()), - }; - const auto code = SM120TF32HCPrenormGemmRuntime::generate(args); - const auto runtime = compiler->build("sm120_tf32_hc_prenorm_gemm", code); - SM120TF32HCPrenormGemmRuntime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("sm120_tf32_hc_prenorm_gemm", std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&sm120_tf32_hc_prenorm_gemm_impl< + {}, {}, + {}, {}, {}, + {}, + {}, + {}, {} + >); +}}; +)", + n, k, + block_m, block_n, block_k, + num_splits, + num_stages, + num_math_threads, num_tma_threads)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(grid_size, 1, 1), + .block_dim = dim3(num_threads, 1, 1), + }, + m, tensor_map_a, tensor_map_b, + reinterpret_cast(d.data_ptr()), reinterpret_cast(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_mega_moe.hpp b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp index b1b408dc18..9073144c1a 100644 --- a/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_mega_moe.hpp @@ -1,133 +1,24 @@ #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/sm90_mega_moe.hpp" +#include "runtime_utils.hpp" namespace deep_gemm { // SM90 uses FP8 operands, float scale factors, register-resident WGMMA // accumulators, and one CTA per scheduled work item. -class SM90FP8MegaMoERuntime final : public LaunchRuntime { -public: - enum class KernelPhase { - Linear1, - Linear2 - }; - - struct Args { - // Templated arguments - int num_max_tokens_per_rank; - int hidden, intermediate_hidden; - int num_experts, num_topk; - int num_ranks; - float activation_clamp; - bool fast_math; - bool bf16_scaled_accum; - KernelPhase kernel_phase; - MegaMoESM90Config config; - - // Runtime arguments - void* y; - int* cumulative_local_expert_recv_stats; - int num_tokens; - layout::SymBuffer<> sym_buffer_ptrs; - - // Tensormaps for activations and weights. Weight scale factors use - // block (128, 128) quantization and are loaded by the math warpgroup - // directly from global memory (no TMA descriptor required). - CUtensorMap tensor_map_l1_acts; - CUtensorMap tensor_map_l1_acts_sf; - CUtensorMap tensor_map_l1_weights; - const float* l1_weights_sf; - CUtensorMap tensor_map_l1_output; - CUtensorMap tensor_map_l2_acts; - CUtensorMap tensor_map_l2_acts_sf; - CUtensorMap tensor_map_l2_weights; - const float* l2_weights_sf; - - // Launch configs - LaunchArgs launch_args; - }; - - static std::string generate_impl(const Args& args) { - const char* kernel_symbol = args.kernel_phase == KernelPhase::Linear1 ? "sm90_fp8_mega_moe_l1_impl" : - "sm90_fp8_mega_moe_l2_impl"; - const auto phase_template_args = args.kernel_phase == KernelPhase::Linear1 ? - fmt::format(",\n {}", args.config.nmajor_schedule ? "true" : "false") : - fmt::format(",\n {}, {}, {}", - args.config.direct_l2_scatter ? "true" : "false", - args.config.nmajor_schedule ? "true" : "false", - args.config.one_warp_cleanup ? "true" : "false"); - return fmt::format(R"( -#include - -using namespace deep_gemm; - -static void __instantiate_kernel() {{ - auto ptr = reinterpret_cast(&{}< - {}, - {}, {}, - {}, {}, - {}, - {}, {}, {}, - {}, - {}, {}, - {}, - {}, {}, {}, - {}, {}, - {}, - {}, - {}, - {}{} - >); -}}; -)", - kernel_symbol, - args.num_max_tokens_per_rank, - args.hidden, args.intermediate_hidden, - args.num_experts, args.num_topk, - args.config.num_experts_per_wave, - args.config.block_m, args.config.block_n, args.config.block_k, - args.config.num_max_pool_tokens, - args.config.num_padded_sf_pool_tokens, - args.config.sf_pool_stride_tokens, - args.config.num_stages, - args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads, - args.config.num_sms, args.num_ranks, - to_string(args.activation_clamp), - args.fast_math ? "true" : "false", - args.config.swap_ab ? "true" : "false", - args.bf16_scaled_accum ? "true" : "false", - phase_template_args); - } - - static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) { - DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config, - args.y, - args.cumulative_local_expert_recv_stats, - args.num_tokens, - args.sym_buffer_ptrs, - args.tensor_map_l1_acts, - args.tensor_map_l1_acts_sf, - args.tensor_map_l1_weights, - args.l1_weights_sf, - args.tensor_map_l1_output, - args.tensor_map_l2_acts, - args.tensor_map_l2_acts_sf, - args.tensor_map_l2_weights, - args.l2_weights_sf - )); - } +enum class SM90MegaMoEKernelPhase { + Linear1, + Linear2 }; static void sm90_fp8_mega_moe( @@ -151,7 +42,7 @@ static void sm90_fp8_mega_moe( // Resolve hardware, generic fallback, phase schedules, and numerical modes // once. The runtime only consumes the resulting complete launch config. - const int num_sms = device_runtime->get_num_sms(); + const int num_sms = runtime->get_num_sms(); const Sm90MoeHeuristicInput heuristic_input { num_sms, num_ranks, num_experts, num_experts_per_rank, @@ -249,55 +140,90 @@ static void sm90_fp8_mega_moe( if (cumulative_local_expert_recv_stats.has_value()) cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr(); - // Launch + // Compile and launch the two phases (Linear1 and Linear2) as separate kernels + // Weight scale factors use block (128, 128) quantization and are loaded by the math + // warpgroup directly from global memory (no TMA descriptor required). const bool bf16_scaled_accum = launch_config.numerical.bf16_scaled_accum; - const SM90FP8MegaMoERuntime::Args args = { - .num_max_tokens_per_rank = num_max_tokens_per_rank, - .hidden = hidden, .intermediate_hidden = intermediate_hidden, - .num_experts = num_experts, .num_topk = num_topk, - .num_ranks = num_ranks, - .activation_clamp = activation_clamp, - .fast_math = fast_math, - .bf16_scaled_accum = bf16_scaled_accum, - .kernel_phase = SM90FP8MegaMoERuntime::KernelPhase::Linear1, - .config = l1_config, - .y = y.data_ptr(), - .cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr, - .num_tokens = num_tokens, - .sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx), - .tensor_map_l1_acts = tensor_map_l1_acts, - .tensor_map_l1_acts_sf = tensor_map_l1_acts_sf, - .tensor_map_l1_weights = tensor_map_l1_weights, - .l1_weights_sf = l1_weights_sf.data_ptr(), - .tensor_map_l1_output = tensor_map_l1_output, - .tensor_map_l2_acts = tensor_map_l2_acts, - .tensor_map_l2_acts_sf = tensor_map_l2_acts_sf, - .tensor_map_l2_weights = tensor_map_l2_weights, - .l2_weights_sf = l2_weights_sf.data_ptr(), - .launch_args = LaunchArgs(l1_config.num_sms, - l1_config.num_dispatch_threads + l1_config.num_non_epilogue_threads + - l1_config.num_epilogue_threads, - l1_config.smem_size, 1) - }; - const auto launch_with_phase = [&](const SM90FP8MegaMoERuntime::KernelPhase kernel_phase, + const auto sym_buffer = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx); + const auto launch_with_phase = [&](const SM90MegaMoEKernelPhase& kernel_phase, const char* kernel_name) { - auto split_args = args; - split_args.kernel_phase = kernel_phase; - const bool is_linear2 = - kernel_phase == SM90FP8MegaMoERuntime::KernelPhase::Linear2; - split_args.config = is_linear2 ? l2_config : l1_config; - split_args.launch_args = LaunchArgs( - split_args.config.num_sms, - split_args.config.num_dispatch_threads + split_args.config.num_non_epilogue_threads + - split_args.config.num_epilogue_threads, - split_args.config.smem_size, 1); - const auto code = SM90FP8MegaMoERuntime::generate(split_args); - const auto runtime = compiler->build(kernel_name, code); - SM90FP8MegaMoERuntime::launch(runtime, split_args); + const bool is_linear2 = kernel_phase == SM90MegaMoEKernelPhase::Linear2; + const auto& config = is_linear2 ? l2_config : l1_config; + const auto phase_template_args = is_linear2 ? + std::format(",\n {}, {}, {}", + config.direct_l2_scatter ? "true" : "false", + config.nmajor_schedule ? "true" : "false", + config.one_warp_cleanup ? "true" : "false") : + std::format(",\n {}", config.nmajor_schedule ? "true" : "false"); + + // Compile + const auto kernel = jit->compile(kernel_name, std::format(R"( +#include + +using namespace deep_gemm; + +static void __instantiate_kernel() {{ + auto ptr = reinterpret_cast(&{}< + {}, + {}, {}, + {}, {}, + {}, + {}, {}, {}, + {}, + {}, {}, + {}, + {}, {}, {}, + {}, {}, + {}, + {}, + {}, + {}{} + >); +}}; +)", + kernel_name, + num_max_tokens_per_rank, + hidden, intermediate_hidden, + num_experts, num_topk, + config.num_experts_per_wave, + config.block_m, config.block_n, config.block_k, + config.num_max_pool_tokens, + config.num_padded_sf_pool_tokens, + config.sf_pool_stride_tokens, + config.num_stages, + config.num_dispatch_threads, config.num_non_epilogue_threads, config.num_epilogue_threads, + config.num_sms, num_ranks, + to_string(activation_clamp), + fast_math ? "true" : "false", + config.swap_ab ? "true" : "false", + bf16_scaled_accum ? "true" : "false", + phase_template_args)); + + // Launch + jit->launch( + kernel, { + .num_smem_bytes = config.smem_size, + .grid_dim = dim3(config.num_sms, 1, 1), + .block_dim = dim3(config.num_dispatch_threads + config.num_non_epilogue_threads + + config.num_epilogue_threads, 1, 1), + }, + y.data_ptr(), + cumulative_local_expert_recv_stats_ptr, + num_tokens, + sym_buffer, + tensor_map_l1_acts, + tensor_map_l1_acts_sf, + tensor_map_l1_weights, + l1_weights_sf.data_ptr(), + tensor_map_l1_output, + tensor_map_l2_acts, + tensor_map_l2_acts_sf, + tensor_map_l2_weights, + l2_weights_sf.data_ptr()); }; - launch_with_phase(SM90FP8MegaMoERuntime::KernelPhase::Linear1, "sm90_fp8_mega_moe_l1_impl"); - launch_with_phase(SM90FP8MegaMoERuntime::KernelPhase::Linear2, "sm90_fp8_mega_moe_l2_impl"); + launch_with_phase(SM90MegaMoEKernelPhase::Linear1, "sm90_fp8_mega_moe_l1_impl"); + launch_with_phase(SM90MegaMoEKernelPhase::Linear2, "sm90_fp8_mega_moe_l2_impl"); } } // 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 f8c720144d..04e4cde73a 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,82 +54,52 @@ 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_clusters; - bool is_varlen; + DG_HOST_ASSERT(128 % num_heads == 0); - int batch_size; - int next_n; - int num_next_n_atoms; - 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_clusters, 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, - static_cast(args.num_next_n_atoms), - 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 + ); +} +// `num_clusters` is the scheduler slot count: one entry per two-CTA cluster for +// next_n=4 multicast, one per SM otherwise (callers pre-divide). static void sm90_paged_mqa_logits_metadata(const torch::Tensor& context_lens, const torch::Tensor& schedule_metadata, const int& batch_size, const int& next_n, @@ -215,103 +116,35 @@ 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_clusters = num_clusters, - .is_varlen = is_varlen, - .batch_size = batch_size, - .next_n = next_n, - .num_next_n_atoms = num_next_n_atoms, - .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; - int num_kv_multicast; - - 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, - args.num_kv_multicast, - 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_clusters, 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, + static_cast(num_next_n_atoms), + 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, @@ -338,7 +171,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(block_kv == 32 or block_kv == 64); DG_HOST_ASSERT(split_kv % mma_m == 0 and logits_stride % split_kv == 0); DG_HOST_ASSERT(not is_varlen); @@ -389,39 +222,51 @@ static void sm90_fp8_paged_mqa_logits(const torch::Tensor& q, smem_umma_barriers + smem_tmem_ptr; DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity); - 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 = 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, - .num_kv_multicast = num_kv_multicast, - .launch_args = LaunchArgs(num_sms, - num_specialized_threads + num_math_threads, - smem_size, num_kv_multicast) - }; - 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, + num_kv_multicast, + to_string(logits_dtype))); + + // Launch: one CTA per SM, grouped into `num_kv_multicast`-CTA clusters + 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), + .cluster_dim = dim3(num_kv_multicast, 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 index ebe4c7a6b9..20f5da312d 100644 --- a/csrc/jit_kernels/impls/smxx_clean_logits.hpp +++ b/csrc/jit_kernels/impls/smxx_clean_logits.hpp @@ -1,33 +1,27 @@ #pragma once -#include "../../jit/compiler.hpp" -#include "../../jit/device_runtime.hpp" -#include "../../jit/kernel_runtime.hpp" +#include + +#include "../../runtime/runtime.hpp" #include "../../utils/exception.hpp" +#include "runtime_utils.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; - }; +// Standalone logits cleaner. The SM100/SM90 MQA kernels fuse cleaning; this is +// kept for the SM120 path, whose kernels do not. +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); - static std::string generate_impl(const Args& args) { - return fmt::format(R"( + // Compile + const auto kernel = jit->compile("smxx_clean_logits", std::format(R"( #include using namespace deep_gemm; @@ -37,45 +31,20 @@ static void __instantiate_kernel() {{ {}, {}, {}, {} >); }}; -)", 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); +)", next_n, block_kv, num_warps, to_string(logits.scalar_type()))); // 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); + jit->launch( + kernel, { + .num_smem_bytes = smem_size, + .grid_dim = dim3(runtime->get_num_sms(), 1, 1), + .block_dim = dim3(num_warps * 32, 1, 1), + }, + seq_len, seq_len_kv, static_cast(stride_logits), + cu_seq_len_k_start.has_value() ? cu_seq_len_k_start.value().data_ptr() : nullptr, + cu_seq_len_k_end.data_ptr(), + logits.data_ptr() + ); } } // 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 82de55e198..fdf0f74df4 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)); - const auto code = TransposeFP32Runtime::generate(args); - const auto runtime = compiler->build("transpose_fp32", code); - TransposeFP32Runtime::launch(runtime, args); + // Compile + const auto kernel = jit->compile("transpose_fp32", std::format(R"( +#include + +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,33 +148,43 @@ 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); - const int max_smem = static_cast(device_runtime->get_prop()->sharedMemPerBlockOptin); - const int block_mn = std::min(48, (max_smem / (sf_k * 4) / 4) * 4); - DG_HOST_ASSERT(block_mn >= 4); + 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; + // Bound `block_mn` by the per-block shared memory capacity (e.g. SM120 has only 99 KB) + const int max_smem = jit->device.get_num_smem_bytes() - psum_smem_elems * 4; + const int block_mn = std::min(48, (max_smem / (smem_sf_k * 4) / 4) * 4); + DG_HOST_ASSERT(block_mn >= 4); + const auto smem_size = block_mn * smem_sf_k * 4 + psum_smem_elems * 4; + + // Compile + const auto kernel = jit->compile("transpose_and_pack_fp32_into_ue8m0", std::format(R"( +#include - const auto code = TransposeAndPackFP32IntoUE8M0Runtime::generate(args); - const auto runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code); - TransposeAndPackFP32IntoUE8M0Runtime::launch(runtime, args); +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) @@ -231,25 +194,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; } @@ -260,8 +224,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()); @@ -280,11 +243,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) @@ -295,7 +262,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, @@ -308,12 +277,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; } @@ -324,9 +292,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()); @@ -340,12 +307,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 55c0fa2b33..5b3bb0d289 100644 --- a/csrc/python_api.cpp +++ b/csrc/python_api.cpp @@ -1,14 +1,19 @@ #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/mega_moe.hpp" #include "apis/sm90_mega.hpp" -#include "apis/runtime.hpp" +#include "apis/mega_mhc.hpp" +#include "apis/mega_gate.hpp" #ifndef TORCH_EXTENSION_NAME #define TORCH_EXTENSION_NAME _C @@ -18,6 +23,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); @@ -26,5 +38,6 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { deep_gemm::layout::register_apis(m); deep_gemm::mega::register_apis(m); deep_gemm::mega::register_sm90_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..219733f7ff --- /dev/null +++ b/csrc/runtime/jit.hpp @@ -0,0 +1,44 @@ +#pragma once + +#include +#include +#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"); + // Allow overriding the in-kernel barrier timeout, see `comm/barrier.cuh` + // NOTES: the flags are part of the kernel cache key, so changing this value + // naturally invalidates the JIT cache + if (const auto barrier_timeout_seconds = deep_jit::get_env("DG_JIT_BARRIER_TIMEOUT_SECONDS", 0); + barrier_timeout_seconds > 0) { + runtime->default_compiler_options.nvcc_flags-> + emplace_back(std::format("-DDG_BARRIER_TIMEOUT_SECONDS={}", barrier_timeout_seconds)); + } + 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 07a81c4e37..136135d6fb 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 48819ca4ee..7c7fce9e11 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 ( @@ -108,27 +113,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 e2ef54a946..713d928699 100644 --- a/deep_gemm/include/deep_gemm/comm/barrier.cuh +++ b/deep_gemm/include/deep_gemm/comm/barrier.cuh @@ -9,8 +9,36 @@ namespace deep_gemm::comm { -// 60s timeout, at 2 GHz -constexpr int64_t kNumTimeoutCycles = 60ll * 2000000000ll; +// Barrier timeout in seconds, overridable at JIT compile time. The host-side JIT +// turns the `DG_JIT_BARRIER_TIMEOUT_SECONDS` environment variable into this +// definition, so slow-but-healthy rendezvous (e.g. large expert-parallel groups +// whose first collective follows a long warm-up) can be distinguished from a +// genuine hang. The flag is part of the kernel signature, so changing the value +// naturally invalidates the JIT cache. +#ifndef DG_BARRIER_TIMEOUT_SECONDS +#define DG_BARRIER_TIMEOUT_SECONDS 60 +#endif + +// Timeout in cycles, at 2 GHz +constexpr int64_t kNumTimeoutCycles = static_cast(DG_BARRIER_TIMEOUT_SECONDS) * 2000000000ll; + +// Spin until `pred()` holds; on timeout, `print_timeout()` runs before the assertion +// (Diagnostic policy), or the kernel traps without printing (TrapOnly policy) +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) { + if constexpr (kTimeoutPolicy == BarrierTimeoutPolicy::Diagnostic) { + print_timeout(); + DG_DEVICE_ASSERT(false and "Timeout"); + } else { + DG_TRAP_ONLY_DEVICE_ASSERT(false and "Timeout"); + } + } + } +} CUTLASS_DEVICE void cluster_sync_with_relaxed_arrive() { // Perform cluster_sync with `barrier.cluster.arrive.relaxed` @@ -33,15 +61,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) { - handle_grid_sync_timeout( - sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, - old_value ^ kFinishSumTag); - } - } 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 (%ds): sm=%u, thread=%u, grid_sync_idx=%u, old=%u, current=%u, expected_tag=%u\n", + DG_BARRIER_TIMEOUT_SECONDS, sm_idx, thread_idx, kGridSyncIndex, old_value, new_value, old_value ^ kFinishSumTag); + }); } sync_scope(); } @@ -60,8 +83,7 @@ CUTLASS_DEVICE void nvlink_barrier(const layout::Workspace& workspace, // Grid sync before NVLink signaling if (sync_prologue) - grid_sync( - workspace, sm_idx, thread_idx, sync_scope); + grid_sync(workspace, sm_idx, thread_idx, sync_scope); // NVLink cross-rank barrier, only SM 0 participates if (sm_idx == 0) { @@ -79,22 +101,16 @@ 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) { - handle_nvlink_barrier_timeout( - sym_buffer.rank_idx, *counter_ptr, - ptx::ld_acq_sys(signal_ptr), target, - signal_phase, signal_sign, kTag); - } - } + wait_until([&]() { return ptx::ld_acq_sys(signal_ptr) == target; }, [&]() { + printf("DeepGEMM NVLink barrier timeout (%ds): rank=%d, counter=%d, signal=%d, target=%d, phase=%d, sign=%d, tag=%d\n", + DG_BARRIER_TIMEOUT_SECONDS, sym_buffer.rank_idx, *counter_ptr, ptx::ld_acq_sys(signal_ptr), target, signal_phase, signal_sign, kTag); + }); } } // Grid sync after NVLink completion if (sync_epilogue) - grid_sync( - workspace, sm_idx, thread_idx, sync_scope); + grid_sync(workspace, sm_idx, thread_idx, sync_scope); } } // namespace deep_gemm::comm 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 45405daa22..0ddb642cf7 100644 --- a/deep_gemm/include/deep_gemm/common/types.cuh +++ b/deep_gemm/include/deep_gemm/common/types.cuh @@ -7,21 +7,12 @@ namespace deep_gemm { enum class MmaKind { BF16 = 0, MXFP8FP4 = 1, + // MXFP4 x MXFP4 with per-32-element UE8M0 scale factors + MXF4 = 2, // NVFP4: both operands are packed E2M1 with per-16-element E4M3 scale factors - NVFP4 = 2, + NVFP4 = 3, }; -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; - // NVFP4 is sub-byte and has no integer byte size. Bit-aware MegaMoE - // helpers handle it explicitly; generic descriptor paths must not. - case MmaKind::NVFP4: return 0; - // Keep unknown kinds distinguishable from the valid sub-byte NVFP4 kind. - default: return -1; - } -} enum class GemmType { Normal = 0, @@ -55,4 +46,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 820aeea6e5..d82ed2946c 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 fe7ede631a..26d5b428fb 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 @@ -38,6 +38,7 @@ template < bool kUseSiTU, float kSiTUBeta, float kSiTULinearBeta, + typename weight_dtype_t, bool kHasShared = (kNumSharedExperts > 0), uint32_t L1_SHAPE_N = kIntermediateHidden * 2, uint32_t L1_SHAPE_K = kHidden, @@ -130,6 +131,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; @@ -144,10 +147,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 @@ -192,7 +196,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]; @@ -274,9 +278,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< @@ -311,8 +316,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 @@ -348,7 +352,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); } @@ -388,6 +392,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; @@ -560,6 +570,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"); @@ -581,10 +596,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) @@ -655,7 +666,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(); @@ -697,17 +708,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); @@ -781,12 +792,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); } @@ -804,7 +817,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 >(); @@ -815,13 +828,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 @@ -866,8 +878,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) { @@ -893,13 +905,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, @@ -1103,8 +1110,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) @@ -1131,12 +1138,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); @@ -1175,10 +1184,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(); } @@ -1208,8 +1215,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( @@ -1241,6 +1247,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 @@ -1299,17 +1317,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]) + @@ -1338,19 +1346,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); @@ -1386,97 +1383,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; + }; - // Cast + // 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); + + // 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..ee061fdc17 --- /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 % 16 == 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/sm120_bf16_gemm.cuh b/deep_gemm/include/deep_gemm/impls/sm120_bf16_gemm.cuh index 31164bfe66..e8c73cdec7 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_bf16_gemm.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_bf16_gemm.cuh @@ -184,8 +184,8 @@ sm120_bf16_gemm_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, const auto a_base = reinterpret_cast(gmem_a_ptr); const auto b_base = reinterpret_cast(gmem_b_ptr); - const uint64_t a_offset = static_cast(scheduler.current_k_cumsum) * shape_m * 2; - const uint64_t b_offset = static_cast(scheduler.current_k_cumsum) * shape_n * 2; + const uint64_t a_offset = static_cast(scheduler.current_k_start) * shape_m * 2; + const uint64_t b_offset = static_cast(scheduler.current_k_start) * shape_n * 2; ptx::tensor_map_replace_global_addr_in_smem(smem_tm_a, a_base + a_offset); ptx::tensor_map_replace_global_addr_in_smem(smem_tm_b, b_base + b_offset); diff --git a/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh b/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh index 4f2d312bf0..6d134bb98b 100644 --- a/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm120_fp8_fp4_gemm_1d1d.cuh @@ -243,22 +243,22 @@ sm120_fp8_fp4_gemm_1d1d_impl(cd_dtype_t* gmem_d, const cd_dtype_t* gmem_c, if constexpr (kKGroupedConstantStride) { const uint64_t a_k_byte_offset = kIsFP4 - ? (static_cast(scheduler.current_k_cumsum) / 2) - : (static_cast(scheduler.current_k_cumsum)); + ? (static_cast(scheduler.current_k_start) / 2) + : (static_cast(scheduler.current_k_start)); const uint64_t b_k_byte_offset = (kIsFP4 || kBIsFP4) - ? (static_cast(scheduler.current_k_cumsum) / 2) - : (static_cast(scheduler.current_k_cumsum)); + ? (static_cast(scheduler.current_k_start) / 2) + : (static_cast(scheduler.current_k_start)); ptx::tensor_map_replace_global_addr_in_smem(smem_tm_a, a_base + a_k_byte_offset); ptx::tensor_map_replace_global_addr_in_smem(smem_tm_b, b_base + b_k_byte_offset); ptx::tensor_map_replace_global_dim_in_smem(smem_tm_a, scheduler.current_shape_k); ptx::tensor_map_replace_global_dim_in_smem(smem_tm_b, scheduler.current_shape_k); } else { const uint64_t a_offset = kIsFP4 - ? (static_cast(scheduler.current_k_cumsum) * shape_m / 2) - : (static_cast(scheduler.current_k_cumsum) * shape_m); + ? (static_cast(scheduler.current_k_start) * shape_m / 2) + : (static_cast(scheduler.current_k_start) * shape_m); const uint64_t b_offset = (kIsFP4 || kBIsFP4) - ? (static_cast(scheduler.current_k_cumsum) * shape_n / 2) - : (static_cast(scheduler.current_k_cumsum) * shape_n); + ? (static_cast(scheduler.current_k_start) * shape_n / 2) + : (static_cast(scheduler.current_k_start) * shape_n); ptx::tensor_map_replace_global_addr_in_smem(smem_tm_a, a_base + a_offset); ptx::tensor_map_replace_global_addr_in_smem(smem_tm_b, b_base + b_offset); const uint64_t a_new_stride = kIsFP4 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_mega_moe.cuh b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh index bf362127ab..8887ddef53 100644 --- a/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh +++ b/deep_gemm/include/deep_gemm/impls/sm90_fp8_mega_moe.cuh @@ -50,15 +50,28 @@ enum class MegaMoEPhaseKind { Linear2 }; +// Power-of-two helpers for the UE8M0 activation SF (previously in `common/math.cuh`) +__forceinline__ __device__ float sm90_fp8_mega_moe_fast_pow2(const int& x) { + uint32_t bits_x = (x + 127) << 23; + return *reinterpret_cast(&bits_x); +} + +__forceinline__ __device__ int sm90_fp8_mega_moe_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); +} + __forceinline__ __device__ void sm90_fp8_mega_moe_get_e4m3_sf_and_sf_inv( const float2& amax, float2& sf, float2& sf_inv) { constexpr float kScale = 1.0f / 448.0f; const auto scaled = make_float2( __fmul_rn(amax.x, kScale), __fmul_rn(amax.y, kScale)); - const auto exp_x = math::fast_log2_ceil(scaled.x); - const auto exp_y = math::fast_log2_ceil(scaled.y); - sf.x = math::fast_pow2(exp_x), sf_inv.x = math::fast_pow2(-exp_x); - sf.y = math::fast_pow2(exp_y), sf_inv.y = math::fast_pow2(-exp_y); + const auto exp_x = sm90_fp8_mega_moe_fast_log2_ceil(scaled.x); + const auto exp_y = sm90_fp8_mega_moe_fast_log2_ceil(scaled.y); + sf.x = sm90_fp8_mega_moe_fast_pow2(exp_x), sf_inv.x = sm90_fp8_mega_moe_fast_pow2(-exp_x); + sf.y = sm90_fp8_mega_moe_fast_pow2(exp_y), sf_inv.y = sm90_fp8_mega_moe_fast_pow2(-exp_y); } template #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_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 f290677d1d..280d3dac86 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,66 @@ 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]; + + // Legacy per-ring-block L2 full counters, used by the kernels that still track L1 -> L2 + // readiness with a count instead of `l2_full_mask` (NVFP4 MegaMoE) + uint32_t l2_full_count[kNumMaxRingBlocks]; +}; +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 +112,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); - - // L2 full block count (ring) - num_bytes += num_ring_blocks * sizeof(uint32_t); + // Fixed-position global and expert signals + num_bytes += sizeof(MegaMoESignals); - // 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 +148,110 @@ 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 + uint64_t* get_l2_full_mask_ptr(const uint32_t& ring_block_idx = 0) const { + return signals->l2_full_mask + ring_block_idx; } + // Legacy count-based L1 -> L2 readiness signal (NVFP4 MegaMoE) 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; + return signals->l2_full_count + 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) @@ -387,10 +410,12 @@ struct MegaMoEBuffer { const auto bf16_token_layout = layout::Data(hidden * 2); const auto intermediate_token_layout = layout::Data(intermediate_hidden * num_mma_elem_bits / 8); const auto shared_intermediate_token_layout = layout::Data(shared_intermediate_hidden * shared_num_mma_elem_bits / 8); - const auto input_sf_layout = layout::Data(with_sf ? hidden / sf_gran_k : 0); - const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / sf_gran_k : 0); + // NOTES: SF records are indexed MN-major by the kernels and never used as a TMA stride, + // so they do not need the 16-byte record alignment (e.g. `intermediate_hidden = 2304`) + const auto input_sf_layout = layout::Data(with_sf ? hidden / sf_gran_k : 0, false); + const auto intermediate_sf_layout = layout::Data(with_sf ? intermediate_hidden / sf_gran_k : 0, false); const auto shared_intermediate_sf_layout = layout::Data( - shared_with_sf ? shared_intermediate_hidden / sf_gran_k : 0); + shared_with_sf ? shared_intermediate_hidden / sf_gran_k : 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); // NVFP4 keeps routing weights in the full-pool tail sidecar because the L1 ring @@ -462,7 +487,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..a94a8a5e71 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)}; @@ -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; @@ -153,6 +158,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 +204,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 +246,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 +281,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 +305,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 +339,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 0348eaeb4f..9fa5e7ad05 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,31 @@ 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)); + } +}; + +// NVFP4: packed E2M1 operands with per-16-element E4M3 scale factors struct SM100_MMA_MXF4NVF4_2x1SM_SS { CUTLASS_DEVICE static void fma(uint64_t const& desc_a, @@ -181,6 +212,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 7bfed53c24..f1442467b3 100644 --- a/deep_gemm/include/deep_gemm/ptx/tma.cuh +++ b/deep_gemm/include/deep_gemm/ptx/tma.cuh @@ -43,6 +43,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" :: @@ -69,7 +86,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" :: @@ -83,7 +100,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), @@ -93,6 +110,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 fccda8d01f..1ad02a516c 100644 --- a/deep_gemm/include/deep_gemm/scheduler/gemm.cuh +++ b/deep_gemm/include/deep_gemm/scheduler/gemm.cuh @@ -34,10 +34,15 @@ template (), - uint32_t kSplitKFactor = 1> + uint32_t kSplitKFactor = 1> // Split-K partitions per MN block (SM120 only) 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 @@ -63,31 +68,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]; } } @@ -97,7 +94,7 @@ struct Scheduler { num_m_blocks = math::ceil_div(shape_m, BLOCK_M); num_n_blocks = math::ceil_div(shape_n, BLOCK_N); num_mn_blocks = num_m_blocks * num_n_blocks; - 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_mn_blocks * kSplitKFactor; } else if constexpr (kGemmType == GemmType::MGroupedContiguous) { @@ -112,13 +109,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(); } } @@ -177,10 +168,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; } @@ -250,23 +238,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 d9bb8f8f79..9533a7920f 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 @@ -404,9 +440,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 eb5b4f4b48..0206b88855 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,110 +63,166 @@ 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) - // An empty varlen batch has a zero-byte dynamic-smem allocation. Handle it - // before touching prefix_work[0], and emit the same one-past-the-end - // sentinel for every scheduler boundary. + // An empty varlen batch may carry no requests at all. Handle it before + // touching shared memory, and emit the same one-past-the-end sentinel + // (q_token_idx = num_q_tokens_total = 0) for every scheduler boundary. if (num_q_tokens_total == 0) { - for (uint32_t sm_idx = thread_idx; sm_idx <= kNumSMs; sm_idx += num_threads) { + for (uint32_t sm_idx = thread_idx; sm_idx <= kNumSMs; sm_idx += kNumMetadataThreads) { schedule_meta[sm_idx * 2] = 0; schedule_meta[sm_idx * 2 + 1] = 0; } return; } - // 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; @@ -257,7 +312,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 3da1297cf4..fa3bd08891 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' or (mma_type == 'fp8xfp4' and activation == 'situ'), \ f'Only FP8xFP4 MegaMoE supports `situ`, got mma_type={mma_type!r}, activation={activation!r}' self.group = group @@ -35,7 +40,7 @@ def __init__(self, group: dist.ProcessGroup, 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, @@ -43,16 +48,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, @@ -117,6 +131,7 @@ def destroy(self): self.group = 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, @@ -125,9 +140,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') @@ -205,7 +217,7 @@ def transform_weights_for_mega_moe( assert activation != 'situ' or (isinstance(l1_weights, tuple) and isinstance(l2_weights, tuple)), \ '`situ` requires FP8xFP4 `(weight, sf)` tuples for both L1 and L2 weights' if isinstance(l1_weights, tuple): - # Scaled FP8xFP4/NVFP4: both formats use the same MN-major SF layout. + # Scaled FP8xFP4/FP8xFP8/NVFP4: all formats use the same MN-major SF layout. # 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])) 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 4dcd7b2d66..80002d092e 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(): @@ -75,10 +79,13 @@ 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: # SM100: FP8_A x FP4_B via UMMA mixed precision quant_config_list.append(QuantConfig((128, 32, False, True))) + quant_config_list.append(QuantConfig((32, 32, True, True))) elif get_arch_major() == 12: quant_config_list.append(QuantConfig((32, 32, True, True))) quant_config_list.append(QuantConfig((128, 32, False, True))) @@ -120,8 +127,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)] @@ -135,31 +143,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: @@ -177,6 +214,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, allow_b_mn): + 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 @@ -199,21 +238,18 @@ def enumerate_m_grouped_masked(dtype: torch.dtype) -> Generator: def enumerate_k_grouped_contiguous(dtype: torch.dtype): arch = get_arch_major() if dtype == torch.bfloat16: - if arch == 9 or arch == 12: - k_alignment_options = [(128, 128)] - else: - k_alignment_options = [(32, 32), (128, 128), (192, 192)] + sf_layout_list = [(128, 128)] if arch in (9, 12) else [(128, 128), (256, 256), (384, 384)] + elif dtype == torch.float4_e2m1fn_x2: + sf_layout_list = [(32, 256), (32, 512), (32, 768)] + elif arch == 9: + sf_layout_list = [(128, 128)] + elif arch == 12: + sf_layout_list = [(32, 128), (128, 128)] else: - if arch == 9: - k_alignment_options = [(128, 128)] - elif arch == 12: - k_alignment_options = [(32, 128), (128, 128)] - else: - k_alignment_options = [(32, 32), (32, 128), (32, 160), (32, 224), - (128, 128), (128, 160), (128, 224)] + sf_layout_list = [(32, 128), (128, 128), (32, 256), (128, 256), (32, 384), (128, 384)] - # SM90 FP8 is K-major, SM120 FP8 supports both NT and TN, all other cases are MN-major. - if arch == 9 and dtype == torch.float8_e4m3fn: + # SM90 FP8 and FP4 are K-major, SM120 FP8 supports both NT and TN, all other cases are MN-major. + if (arch == 9 and dtype == torch.float8_e4m3fn) or dtype == torch.float4_e2m1fn_x2: major_pairs = [(MajorTypeAB.KMajor, MajorTypeAB.KMajor)] elif arch == 12 and dtype == torch.float8_e4m3fn: major_pairs = [(MajorTypeAB.KMajor, MajorTypeAB.KMajor), (MajorTypeAB.MNMajor, MajorTypeAB.MNMajor)] @@ -221,37 +257,37 @@ def enumerate_k_grouped_contiguous(dtype: torch.dtype): major_pairs = [(MajorTypeAB.MNMajor, MajorTypeAB.MNMajor)] psum_list = (False, True) if arch == 10 else (False, ) - # Must use FP32 accumulation and 1D1D kernels. + # SM90 and SM120 kernels only support FP32 accumulation + if arch in (9, 12): + cd_options = [(True, torch.float)] + else: + cd_options = [(True, torch.float), (False, torch.float), (False, torch.bfloat16)] + for major_a, major_b in major_pairs: - # The first shape has many small groups, stressing the SM90 in-place tensor-map update. + # 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 ( 8, 4096, 7168, 4096), ( 8, 7168, 2048, 4096), # EP32 (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]): + yield list(real_ks_cpu) -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)) + 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 - 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 + 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(): @@ -265,26 +301,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(): @@ -331,13 +369,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 @@ -368,6 +408,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)): @@ -378,6 +419,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() @@ -385,20 +427,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=not quant_config.is_fp4_a) - 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] @@ -406,6 +449,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, @@ -421,9 +476,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) @@ -433,81 +494,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 @@ -516,44 +566,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 7a47d478ee..b96cf89dd3 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 @@ -85,6 +85,15 @@ def dtype_tag(dtype: torch.dtype) -> str: return '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] @@ -142,13 +151,13 @@ def enumerate_mqa_logits(): # FP8 uses a per-KV float scale. MXFP4/MXFP8 use packed per-32 block scales. if arch_major == 10: fmts = ('mxfp4', 'mxfp8', 'fp8') - shapes = ((510, 130560), (512, 130560), (2048, 8192), (8192, 65536)) + shapes = ((510, 130560), (512, 130560), (2048, 8192), (2048, 65536), (8192, 8192), (8192, 65536)) elif arch_major == 12: fmts = ('mxfp4', 'fp8') shapes = ((128, 4096), (512, 8192), (2048, 8192), (4096, 8192)) else: fmts = ('fp8', ) - shapes = ((2048, 8192), (8192, 65536)) + shapes = ((2048, 8192), (2048, 65536), (8192, 8192), (8192, 65536)) for fmt in fmts: is_mxfp4 = fmt == 'mxfp4' @@ -164,8 +173,11 @@ def enumerate_mqa_logits(): if weights_dtype == torch.float16 and seq_len % 4 != 0: continue head_dims = (128, ) if arch_major == 12 and is_mxfp4 else ((64, 128) if is_mxfp4 else (32, 64, 128)) - heads = (8, 16, 32, 64) if arch_major == 10 else ((16, 32, 64) if arch_major == 12 else (32, 64)) + heads = (8, 12, 16, 20, 32, 64) if arch_major == 10 else ((16, 32, 64) if arch_major == 12 else (32, 64)) for num_heads in heads: + # The FP16-weights two-CTA kernel tiles Q in `128 / num_heads` rows + if weights_dtype == torch.float16 and 128 % num_heads != 0: + continue for head_dim in head_dims: for disable_cp in (False, True): if not disable_cp and (seq_len_kv % seq_len != 0 or seq_len % 2 != 0): @@ -182,7 +194,7 @@ def enumerate_mqa_logits(): weights = torch.randn(seq_len, num_heads, device='cuda', dtype=torch.float32) # FP16 weights select nv_dev's SM100-only two-CTA accumulator kernel. Scale # them down to avoid overflowing its FP16 score/reduction intermediates. - kernel_weights = (weights * 0.1).to(weights_dtype) if weights_dtype == torch.float16 else weights.to(weights_dtype) + kernel_weights = to_mqa_weights(weights * 0.1 if weights_dtype == torch.float16 else weights, weights_dtype) ks, ke = generate_ks_ke_tests(seq_len, seq_len_kv, disable_cp) if compressed_logits and weights_dtype == torch.float16: # Adjacent rows deliberately use disjoint windows. The FP16 kernel @@ -244,6 +256,15 @@ 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 + # The FP16-weights kernel (nv_dev) has no scheduled-metadata path + if get_arch_major() == 10 and weights_dtype != torch.float16: + 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) @@ -268,17 +289,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() @@ -315,6 +340,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 @@ -331,28 +384,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() # Varlen is SM100/SM120-only (SM90 kernel statically rejects it). SM90 supports @@ -378,8 +409,8 @@ def enumerate_paged_mqa_logits(): for batch_size in (256, 4096): next_ns = (1, ) if is_varlen else ((1, 2, 4, 5, 6) if arch_major == 10 else ((1, 2, 3, 4, 5, 6) if arch_major == 12 else (1, 2, 4))) for next_n in next_ns: - for max_tokens_per_batch in ((1, 4, 10) if is_varlen else (1, )): - heads = (8, 16, 32, 64) if arch_major == 10 else ((16, 32, 64) if arch_major == 12 else (32, 64)) + for max_tokens_per_batch in ((1, 4, 6, 10) if is_varlen else (1, )): + heads = (8, 12, 16, 20, 32, 64) if arch_major == 10 else ((16, 32, 64) if arch_major == 12 else (32, 64)) if is_mxfp4: head_dims = (128, ) if arch_major == 12 else (64, 128) else: @@ -428,7 +459,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) # SM90 consumes two 32-token physical pages per 64-token MMA. Keep one # deterministic case at an exact three-page table width so the final MMA @@ -522,10 +553,14 @@ def enumerate_paged_mqa_logits(): # so the metadata schedule must be sized for clusters, not SMs. num_kv_multicast = 2 if get_arch_major() == 9 and next_n == 4 else 1 num_clusters = deep_gemm.get_num_sms() // num_kv_multicast + metadata_kwargs = dict( + context_lens=context_lens_nextn, block_kv=block_kv, + num_sms=num_clusters, 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, num_clusters, 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, ) @@ -560,16 +595,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 @@ -579,6 +618,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__': @@ -588,3 +895,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 35772d5465..b1b4ba3351 100644 --- a/tests/test_bf16.py +++ b/tests/test_bf16.py @@ -8,46 +8,52 @@ 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 uses different kernels by shape: nvjet/cutlass *gemm* for general - # shapes, a dot_kernel + reduce_1Block_kernel GEMV path for M=1, plus a + # shapes, a dot_kernel + reduce_1Block_kernel GEMV path for M=1 (SM120), plus a # split-K reduce kernel. Sum all matches for the total cuBLAS time. - cublas_t = sum(bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a, b, d, c=c), - ('nvjet', 'gemm', 'gemv', 'dot', 'reduce'), suppress_kineto_output=True)) + deep_gemm.use_deterministic_algorithms(False) + cublas_t = sum(bench_kineto(lambda: deep_gemm.bf16_gemm_nt(a, b, d, c=c), + ('nvjet', 'gemm', 'gemv', 'dot', 'reduce'), + suppress_kineto_output=True, with_multiple_kernels=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 | ' @@ -70,8 +76,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() @@ -79,20 +85,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(): @@ -122,40 +120,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}): ' @@ -168,47 +156,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() @@ -216,18 +211,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=})') # cuBLAS uses different kernels by shape: nvjet/cutlass *gemm* for general # shapes, and a dot_kernel + reduce_1Block_kernel GEMV path for M=1. Match all. @@ -241,6 +242,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) @@ -255,3 +335,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 f2d391c7a3..8123dd0d0e 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 ) @@ -52,12 +52,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_bf16_einsum_close('bhr,hdr->bhd', z, ref_z, x, y) - 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_bf16_einsum_close('bhr,hdr->bhd', z, ref_z, x, y) + 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 | ' @@ -74,12 +78,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_bf16_einsum_close('bhd,hdr->bhr', z, ref_z, x, y) - 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_bf16_einsum_close('bhd,hdr->bhr', z, ref_z, x, y) + 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 | ' @@ -88,8 +97,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) @@ -103,12 +134,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 | ' @@ -119,6 +174,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): @@ -138,7 +194,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 | ' @@ -162,10 +218,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 | ' @@ -183,6 +239,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 0b5f297c53..89e24f35cd 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, MajorTypeAB, get_ue8m0_usage, layout_masked_to_psum, align, + KernelType, MajorTypeAB, 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' @@ -41,37 +42,57 @@ def test_gemm() -> None: continue 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) - a_k = a[0].contiguous() if not major_a.is_k_major() else a[0] - b_k = b[0].contiguous() if not major_b.is_k_major() else b[0] - # On SM120 (consumer Blackwell) cuBLAS picks either an nvjet GEMM kernel or - # falls back to a legacy sm89_xmma FP8 kernel depending on shape; match both, - # plus the split-K reduce kernel. (SM90/SM100 use nvjet and never the sm89 - # fallback, so 'xmma' simply matches nothing there.) - cublas_nvjet, cublas_xmma, split_k_t = \ - bench_kineto(lambda: deep_gemm.cublaslt_gemm_nt(a_k, b_k, d, c=c), ('nvjet', 'xmma', 'reduce'), suppress_kineto_output=True) \ - if not quant_config.is_fp4_a and not quant_config.is_fp4_b else (0, 0, 0) - cublas_t = cublas_nvjet + cublas_xmma + 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) + # On SM120 (consumer Blackwell) cuBLAS picks either an nvjet GEMM kernel or falls back + # to a legacy sm89_xmma FP8 kernel depending on shape; match both, plus the split-K + # reduce kernel. (SM90/SM100 use nvjet and never the sm89 fallback, so 'xmma' simply + # matches nothing there.) + cublas_times = bench_kineto(cublas_func, ('nvjet', 'bstensorop', 'xmma', '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: @@ -90,30 +111,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) @@ -145,42 +172,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 @@ -197,54 +223,145 @@ 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, _, gran_k, k_alignment, use_psum_layout in enumerate_k_grouped_contiguous(torch.float8_e4m3fn): - k_grouped_fp8_gemm_contiguous = deep_gemm.k_grouped_fp8_gemm_nt_contiguous if major_a == MajorTypeAB.KMajor \ - else deep_gemm.k_grouped_fp8_gemm_tn_contiguous - 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() + # FP8: the K-major layout (SM90 always, SM120 optionally) uses the NT kernel, MN-major uses TN + test_options = [(torch.float8_e4m3fn, QuantConfig(), None)] + 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, fixed_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) + gemm = fixed_gemm + if gemm is None: + gemm = deep_gemm.k_grouped_fp8_gemm_nt_contiguous if major_a == MajorTypeAB.KMajor \ + else deep_gemm.k_grouped_fp8_gemm_tn_contiguous + + 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