diff --git a/.gitignore b/.gitignore index 1a89e54605..a19983f36a 100644 --- a/.gitignore +++ b/.gitignore @@ -13,6 +13,8 @@ flashinfer/_build_meta.py flashinfer/data/ flashinfer/jit/aot_config.py csrc/aot_default_additional_params.h +flashinfer/trtllm/gemm/trtllmGen_gemm_export +flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export # DS_Store files .DS_store diff --git a/csrc/trtllm_batched_gemm_runner.cu b/csrc/trtllm_batched_gemm_runner.cu index b0c43ea751..c4cce95011 100644 --- a/csrc/trtllm_batched_gemm_runner.cu +++ b/csrc/trtllm_batched_gemm_runner.cu @@ -214,6 +214,9 @@ void TrtllmGenBatchedGemmRunner::run( gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m; gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n; gemmData.mProblemDimensions.mK = k; + gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM; + gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN; + gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK; gemmData.mProblemDimensions.mRank = 0; gemmData.mProblemDimensions.mWorldSize = 1; @@ -252,10 +255,6 @@ void TrtllmGenBatchedGemmRunner::run( int32_t multiProcessorCount; cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount, device); - gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM; - gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN; - gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK; - // FIXME once we start using all-reduce in the epilogue of the bmm this can be moved elsewhere bmm.runInitBeforeWorldSync(config, gemmData, static_cast(stream)); @@ -449,6 +448,9 @@ bool TrtllmGenBatchedGemmRunner::isValidConfigIndex(int32_t configIndex, int32_t gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m; gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n; gemmData.mProblemDimensions.mK = k; + gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM; + gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN; + gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK; gemmData.mProblemDimensions.mRank = 0; gemmData.mProblemDimensions.mWorldSize = 1; gemmData.mProblemDimensions.mMaxNumCtasInTokenDim = maxNumCtasInBatchDim; diff --git a/csrc/trtllm_gemm_runner.cu b/csrc/trtllm_gemm_runner.cu index ffe0bc2cf8..7ab52ca06a 100644 --- a/csrc/trtllm_gemm_runner.cu +++ b/csrc/trtllm_gemm_runner.cu @@ -119,6 +119,10 @@ class TrtllmGenGemmRunner { gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m; gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n; gemmData.mProblemDimensions.mK = k; + // TODO(jimmyzho) disable until fix trtllm-gen + // gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM; + // gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN; + // gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK; gemmData.mProblemDimensions.mRank = 0; gemmData.mProblemDimensions.mWorldSize = 1; @@ -139,6 +143,10 @@ class TrtllmGenGemmRunner { gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m; gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n; gemmData.mProblemDimensions.mK = k; + // TODO(jimmyzho) disable until fix trtllm-gen + // gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM; + // gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN; + // gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK; gemmData.mProblemDimensions.mRank = 0; gemmData.mProblemDimensions.mWorldSize = 1; @@ -187,6 +195,10 @@ class TrtllmGenGemmRunner { gemmData.mProblemDimensions.mM = mOptions.transposeMmaOutput ? n : m; gemmData.mProblemDimensions.mN = mOptions.transposeMmaOutput ? m : n; gemmData.mProblemDimensions.mK = k; + // TODO(jimmyzho) disable until fix trtllm-gen + // gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM; + // gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN; + // gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK; gemmData.mProblemDimensions.mRank = 0; gemmData.mProblemDimensions.mWorldSize = 1; diff --git a/csrc/trtllm_low_latency_gemm_runner.cu b/csrc/trtllm_low_latency_gemm_runner.cu index f3ce0d43c3..6b47d2f7cc 100644 --- a/csrc/trtllm_low_latency_gemm_runner.cu +++ b/csrc/trtllm_low_latency_gemm_runner.cu @@ -49,6 +49,10 @@ gemm::gemm::GemmData createGemmData(int64_t m, int64_t n, int64_t k) { gemmData.mProblemDimensions.mM = n; gemmData.mProblemDimensions.mN = m; gemmData.mProblemDimensions.mK = k; + // TODO(jimmyzho) disable until fix trtllm-gen + // gemmData.mProblemDimensions.mValidM = gemmData.mProblemDimensions.mM; + // gemmData.mProblemDimensions.mValidN = gemmData.mProblemDimensions.mN; + // gemmData.mProblemDimensions.mValidK = gemmData.mProblemDimensions.mK; gemmData.mProblemDimensions.mRank = 0; gemmData.mProblemDimensions.mWorldSize = 1; diff --git a/flashinfer/artifacts.py b/flashinfer/artifacts.py index cce73a6827..ebfabfe984 100644 --- a/flashinfer/artifacts.py +++ b/flashinfer/artifacts.py @@ -79,6 +79,54 @@ def get_available_cubin_files( return tuple() +def get_available_header_files( + source: str, retries: int = 3, delay: int = 5, timeout: int = 10 +) -> tuple[str, ...]: + """ + Recursively navigates through child directories (e.g., include/) and finds + all *.h header files, returning them as a tuple of relative paths. + """ + result: list[str] = [] + + def fetch_directory(url: str, prefix: str = "") -> None: + for attempt in range(1, retries + 1): + try: + response = requests.get(url, timeout=timeout) + response.raise_for_status() + + # Find all .h header files in this directory + header_hrefs = re.findall(r'', response.text) + for h in header_hrefs: + result.append(prefix + h if prefix else h) + + # Find all subdirectories (links ending with /) + dir_hrefs = re.findall(r'', response.text) + for d in dir_hrefs: + # Skip parent directory links + if d == "../" or d.startswith(".."): + continue + subdir_url = safe_urljoin(url, d) + subdir_prefix = prefix + d if prefix else d + fetch_directory(subdir_url, subdir_prefix) + + return # Success, exit retry loop + + except requests.exceptions.RequestException as e: + logger.warning( + f"Fetching available header files {url}: attempt {attempt} failed: {e}" + ) + + if attempt < retries: + logger.info(f"Retrying in {delay} seconds...") + time.sleep(delay) + + logger.error(f"Max retries reached for {url}. Fetch failed.") + + fetch_directory(source) + logger.info(f"result: {result}") + return tuple(result) + + @dataclass(frozen=True) class ArtifactPath: """ @@ -89,7 +137,7 @@ class ArtifactPath: TRTLLM_GEN_FMHA: str = "e86f0e45764555d070c3d143b4caaea61a45b777/fmha/trtllm-gen/" TRTLLM_GEN_BMM: str = ( - "456b1ae890d436c794b17e4435b41b849d3e5950/batched_gemm-2a674db-3a84a12" + "de418500652abd6325552396fc931d5165f12804/batched_gemm-2a674db-3a84a12/" ) TRTLLM_GEN_GEMM: str = ( "1fddc48b7b48af33914d040051b3e2ee9ba4701e/gemm-145d1b1-9b113e3" @@ -110,7 +158,7 @@ class CheckSumHash: "c4c93904a4c72b8a3d0d5c525c6decb71c835b477d7d75651ecaaa7007c5a3ef" ) TRTLLM_GEN_BMM: str = ( - "b9121fed5dd7700b7c2a0dcbcf2ef022483855cf585263324275b0072cca6bb7" + "69ade639fa6640ea90261599581c8cda50f8f6b9ac9be07a3ca38f3bf5458452" ) DEEPGEMM: str = "1a2a166839042dbd2a57f48051c82cd1ad032815927c753db269a4ed10d0ffbf" TRTLLM_GEN_GEMM: str = ( @@ -182,6 +230,9 @@ def get_subdir_file_list() -> Generator[tuple[str, str], None, None]: yield (checksum_path, CheckSumHash.map_checksums[checksum_path]) for name in get_available_cubin_files(safe_urljoin(base, cubin_dir)): yield (safe_urljoin(cubin_dir, name), checksums[name]) + for name in get_available_header_files(safe_urljoin(base, cubin_dir)): + full_path = safe_urljoin(cubin_dir, name) + yield (full_path, checksums[full_path]) def download_artifacts() -> None: @@ -190,7 +241,7 @@ def download_artifacts() -> None: # use a shared session to make use of HTTP keep-alive and reuse of # HTTPS connections. session = requests.Session() - cubin_files = list(get_subdir_file_list()) + cubin_files = list[tuple[str, str]](get_subdir_file_list()) num_threads = int(os.environ.get("FLASHINFER_CUBIN_DOWNLOAD_THREADS", "4")) with tqdm_logging_redirect( total=len(cubin_files), desc="Downloading cubins" diff --git a/flashinfer/jit/cubin_loader.py b/flashinfer/jit/cubin_loader.py index 1aae47722a..922afdcb1c 100644 --- a/flashinfer/jit/cubin_loader.py +++ b/flashinfer/jit/cubin_loader.py @@ -17,13 +17,16 @@ import ctypes import hashlib import os +import pathlib from urllib.parse import urljoin import shutil import time +from typing import Union import uuid import filelock +from .utils import write_if_different from .core import logger from .env import FLASHINFER_CUBIN_DIR @@ -136,14 +139,20 @@ def download_file( return False -def get_meta_hash(checksums_bytes: bytes) -> str: +def get_meta_hash( + checksums_bytes: bytes, target_file: str = "flashinferMetaInfo.h" +) -> str: """ Parse the checksums.txt file and get the hash of corresponding flashinferMetaInfo.h file """ checksums_lines = checksums_bytes.decode("utf-8").splitlines() for line in checksums_lines: sha256, filename = line.strip().split() - if ".h" in filename: + # Match on path segment boundary to avoid substring collisions + # (e.g. "Enums.h" must not match "BatchedGemmEnums.h") + if filename.lower() == target_file.lower() or filename.lower().endswith( + "/" + target_file.lower() + ): return sha256 raise ValueError("Invalid checksums.txt, no flashinferMetaInfo.h found") @@ -189,6 +198,27 @@ def load_cubin(cubin_path: str, sha256: str) -> bytes: return b"" +def get_file( + uri_path: str, + sha256: str, + file_path: str, + session=None, +) -> bytes: + """ + Load a file from local cache directory {file_path}, ensure that the sha256 signature matches. + Otherwise, download the file from {uri_path} and write to {file_path}. + """ + + file = load_cubin(file_path, sha256) + if file: + return file + os.makedirs(os.path.dirname(file_path), exist_ok=True) + uri = safe_urljoin(FLASHINFER_CUBINS_REPOSITORY, uri_path) + logger.info(f"Fetching file from {uri}") + download_file(uri, file_path, session=session) + return load_cubin(file_path, sha256) + + def get_cubin(file_name: str, sha256: str, session=None) -> bytes: """ Load a cubin from the local cache directory with {file_name} and @@ -211,6 +241,80 @@ def get_cubin(file_name: str, sha256: str, session=None) -> bytes: return load_cubin(cubin_path, sha256) +def download_trtllm_headers( + op: str, + header_dest_dir: Union[str, pathlib.Path], + header_path: str, + artifact_path: str, + checksum: bytes, +): + header_dest_dir = pathlib.Path(header_dest_dir) + + if op == "bmm": + header_files = [ + "BatchedGemmEnums.h", + "BatchedGemmInterface.h", + "BatchedGemmOptions.h", + "Enums.h", + "GemmGatedActOptions.h", + "GemmOptions.h", + "KernelParams.h", + "KernelParamsDecl.h", + "KernelTraits.h", + "TmaDescriptor.h", + "trtllm/gen/CommonUtils.h", + "trtllm/gen/CudaArchDecl.h", + "trtllm/gen/CudaKernelLauncher.h", + "trtllm/gen/DtypeDecl.h", + "trtllm/gen/MmaDecl.h", + "trtllm/gen/SfLayoutDecl.h", + "trtllm/gen/SparsityDecl.h", + ] + + else: + header_files = [ + "GemmInterface.h", + "GemmOptions.h", + "Enums.h", + "KernelTraits.h", + "KernelParams.h", + "KernelParamsDecl.h", + "TmaDescriptor.h", + "trtllm/gen/CommonUtils.h", + "trtllm/gen/CudaKernelLauncher.h", + "trtllm/gen/DtypeDecl.h", + "trtllm/gen/MmaDecl.h", + "trtllm/gen/SfLayoutDecl.h", + "trtllm/gen/CudaArchDecl.h", + ] + + artifact_hash_path = header_dest_dir / ".artifact_hash" + + # Check if cached headers are from a different artifact version (e.g. after git checkout) + if artifact_hash_path.exists(): + with open(artifact_hash_path, "r") as f: + cached_hash = f.read().strip() + if cached_hash != artifact_path: + raise RuntimeError( + f"Detected inconsistent cached artifacts. " + f"(Cached trtllm headers were downloaded for artifact " + f"'{cached_hash}', but current code expects " + f"'{artifact_path}'). " + f"Please clear the cache to confirm and allow the new headers to be downloaded: " + f"rm -rf {header_dest_dir}." + ) + + for file in header_files: + uri_path = f"{header_path}/{file}" + file_hash = get_meta_hash(checksum, file) + file_path = str(header_dest_dir / file) + result = get_file(uri_path, file_hash, file_path) + assert result, f"{file} not found" + + # Record which artifact version these headers belong to + write_if_different(artifact_hash_path, artifact_path) + + def convert_to_ctypes_char_p(data: bytes): return ctypes.c_char_p(data) diff --git a/flashinfer/jit/fused_moe.py b/flashinfer/jit/fused_moe.py index 1e27515b69..a3240f7353 100644 --- a/flashinfer/jit/fused_moe.py +++ b/flashinfer/jit/fused_moe.py @@ -26,7 +26,7 @@ sm89_nvcc_flags, ) from .cpp_ext import is_cuda_version_at_least -from .cubin_loader import get_cubin, get_meta_hash +from .cubin_loader import get_cubin, get_meta_hash, download_trtllm_headers from .gemm.cutlass.generate_kernels import generate_gemm_operations @@ -233,6 +233,19 @@ def gen_trtllm_gen_fused_moe_sm100_module() -> JitSpec: # make sure "flashinferMetaInfo.h" is downloaded or cached assert metainfo, f"{header_name}.h not found" + header_path = f"{include_path}/trtllmGen_bmm_export" + header_dest_dir = ( + jit_env.FLASHINFER_CUBIN_DIR + / "flashinfer" + / "trtllm" + / "batched_gemm" + / "trtllmGen_bmm_export" + ) + + download_trtllm_headers( + "bmm", header_dest_dir, header_path, ArtifactPath.TRTLLM_GEN_BMM, checksum + ) + # currently only support Blackwell nvcc_flags = current_compilation_context.get_nvcc_flags_list( supported_major_versions=[10] @@ -265,7 +278,7 @@ def gen_trtllm_gen_fused_moe_sm100_module() -> JitSpec: ] + nvcc_flags, extra_include_paths=[ - # link "include" sub-directory in cache + jit_env.FLASHINFER_CUBIN_DIR, jit_env.FLASHINFER_CUBIN_DIR / include_path, jit_env.FLASHINFER_CSRC_DIR / "nv_internal", jit_env.FLASHINFER_CSRC_DIR / "nv_internal/include", diff --git a/flashinfer/jit/gemm/core.py b/flashinfer/jit/gemm/core.py index 335a0c2b48..a170c83128 100644 --- a/flashinfer/jit/gemm/core.py +++ b/flashinfer/jit/gemm/core.py @@ -30,7 +30,11 @@ sm100f_nvcc_flags, current_compilation_context, ) -from ..cubin_loader import get_cubin, get_meta_hash +from ..cubin_loader import ( + # download_trtllm_headers, + get_cubin, + get_meta_hash, +) from ..utils import dtype_cutlass_map, filename_safe_dtype_map, write_if_different @@ -536,6 +540,20 @@ def gen_trtllm_gen_gemm_module() -> JitSpec: ) # make sure "flashinferMetaInfo.h" is downloaded or cached assert metainfo, f"{header_name}.h not found" + + # TODO(jimmyzho): Re-enable after fixing trtllm-gen cubin generation issues. + # header_path = f"{include_path}/trtllmGen_gemm_export" + # header_dest_dir = ( + # jit_env.FLASHINFER_CUBIN_DIR + # / "flashinfer" + # / "trtllm" + # / "gemm" + # / "trtllmGen_gemm_export" + # ) + # download_trtllm_headers( + # "gemm", header_dest_dir, header_path, ArtifactPath.TRTLLM_GEN_GEMM, checksum + # ) + return gen_jit_spec( "trtllm_gemm", [ @@ -549,7 +567,11 @@ def gen_trtllm_gen_gemm_module() -> JitSpec: ] + sm100a_nvcc_flags, # link "include" sub-directory in cache - extra_include_paths=[jit_env.FLASHINFER_CUBIN_DIR / include_path], + extra_include_paths=[ + jit_env.FLASHINFER_CUBIN_DIR / include_path, + # jit_env.FLASHINFER_CUBIN_DIR, + # jit_env.FLASHINFER_CUBIN_DIR / include_path, + ], ) @@ -687,6 +709,20 @@ def gen_trtllm_low_latency_gemm_module() -> JitSpec: ) # make sure "flashinferMetaInfo.h" is downloaded or cached assert metainfo, f"{header_name}.h not found" + + # TODO(jimmyzho): Re-enable after fixing trtllm-gen cubin generation issues. + # header_path = f"{include_path}/trtllmGen_gemm_export" + # header_dest_dir = ( + # jit_env.FLASHINFER_CUBIN_DIR + # / "flashinfer" + # / "trtllm" + # / "gemm" + # / "trtllmGen_gemm_export" + # ) + # download_trtllm_headers( + # "gemm", header_dest_dir, header_path, ArtifactPath.TRTLLM_GEN_GEMM, checksum + # ) + return gen_jit_spec( "trtllm_low_latency_gemm", [ @@ -699,7 +735,10 @@ def gen_trtllm_low_latency_gemm_module() -> JitSpec: f'-DTLLM_GEN_GEMM_CUBIN_PATH=\\"{ArtifactPath.TRTLLM_GEN_GEMM}\\"', ] + sm100a_nvcc_flags, - # link "include" sub-directory in cache - extra_include_paths=[jit_env.FLASHINFER_CUBIN_DIR / include_path], + extra_include_paths=[ + jit_env.FLASHINFER_CUBIN_DIR / include_path + # jit_env.FLASHINFER_CUBIN_DIR, + # jit_env.FLASHINFER_CUBIN_DIR / include_path, + ], extra_ldflags=["-lcuda"], ) diff --git a/flashinfer/jit/moe_utils.py b/flashinfer/jit/moe_utils.py index 71f54eb907..b877ad5e8f 100644 --- a/flashinfer/jit/moe_utils.py +++ b/flashinfer/jit/moe_utils.py @@ -21,6 +21,9 @@ current_compilation_context, ) +from .cubin_loader import download_trtllm_headers, get_cubin +from ..artifacts import ArtifactPath, CheckSumHash + def gen_moe_utils_module() -> JitSpec: """ @@ -33,7 +36,22 @@ def gen_moe_utils_module() -> JitSpec: - moeActivation: Apply activation functions with optional FP4 quantization - moeSort: Sort tokens by expert assignment (DeepSeekV3 routing) """ + + download_trtllm_headers( + "bmm", + jit_env.FLASHINFER_CUBIN_DIR + / "flashinfer" + / "trtllm" + / "batched_gemm" + / "trtllmGen_bmm_export", + f"{ArtifactPath.TRTLLM_GEN_BMM}/include/trtllmGen_bmm_export", + ArtifactPath.TRTLLM_GEN_BMM, + get_cubin( + f"{ArtifactPath.TRTLLM_GEN_BMM}/checksums.txt", CheckSumHash.TRTLLM_GEN_BMM + ), + ) nvcc_flags = [ + "-DTLLM_GEN_EXPORT_INTERFACE", # Use relative includes in downloaded headers "-DENABLE_BF16", "-DENABLE_FP8", "-DENABLE_FP4", @@ -78,13 +96,7 @@ def gen_moe_utils_module() -> JitSpec: / "tensorrt_llm" / "kernels" / "cutlass_kernels", - # Include paths for routing kernels - jit_env.FLASHINFER_INCLUDE_DIR, - # Include path for trtllm/gen/MmaDecl.h (used by DtypeDecl.h) - jit_env.FLASHINFER_INCLUDE_DIR - / "flashinfer" - / "trtllm" - / "batched_gemm" - / "trtllmGen_bmm_export", + # Include paths for routing kernels and downloaded headers + jit_env.FLASHINFER_CUBIN_DIR, ], ) diff --git a/include/flashinfer/trtllm/batched_gemm/KernelRunner.h b/include/flashinfer/trtllm/batched_gemm/KernelRunner.h index 54cd824c0e..64e958d9e8 100644 --- a/include/flashinfer/trtllm/batched_gemm/KernelRunner.h +++ b/include/flashinfer/trtllm/batched_gemm/KernelRunner.h @@ -21,8 +21,8 @@ #include #include -#include "trtllmGen_bmm_export/Enums.h" -#include "trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h" +#include "flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h" namespace tensorrt_llm { namespace kernels { diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h deleted file mode 100644 index 919d6cb00d..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmEnums.h +++ /dev/null @@ -1,61 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include - -namespace batchedGemm { - -namespace batchedGemm { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -enum class RouteImpl { - // No Routing - NoRoute = 0, - // Use LDGSTS to do the routing - Ldgsts = 1, - // Use UTMALDG.GATHER4 to do the routing - Tma = 2, - // Use LDG+STS to do the routing - LdgPlusSts = 3 -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool doesRouteImplUseNoRoute(RouteImpl mode) { return (mode == RouteImpl::NoRoute); } - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool doesRouteImplUseLdgsts(RouteImpl mode) { return (mode == RouteImpl::Ldgsts); } - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool doesRouteImplUseTma(RouteImpl mode) { return (mode == RouteImpl::Tma); } - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool doesRouteImplUseLdgPlusSts(RouteImpl mode) { return (mode == RouteImpl::LdgPlusSts); } - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace batchedGemm - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h deleted file mode 100644 index 99e95d760a..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmInterface.h +++ /dev/null @@ -1,876 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include -#include - -#include "BatchedGemmOptions.h" -#include "KernelParams.h" -#include "trtllm/gen/CudaKernelLauncher.h" - -#ifdef TLLM_GEN_EXPORT_INTERFACE -#ifdef TLLM_GEN_EXPORT_FLASHINFER -#include "flashinferMetaInfo.h" -#else -#include "KernelMetaInfo.h" -#endif // TLLM_GEN_EXPORT_FLASHINFER -#endif // TLLM_GEN_EXPORT_INTERFACE - -namespace batchedGemm { - -namespace batchedGemm { - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// BatchedGemmData -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -struct BatchedGemmData { - struct ProblemDimensions { - // The number of batches. - int32_t mNumBatches{0}; - // The number of tokens. Set to 0 if routeAct is false. - int32_t mNumTokens{0}; - // Whether the batch is on the M dimension. - bool mBatchM{true}; - // The maximum number of CTAs in the token dimension. - // Need to be set if mNumTokens > 0 and the token per batch - // distribution is not known at launch time. - // In this case, the kernel will launch mMaxNumCtasInTokenDim CTAs in token dim and exit early - // if the idx of CTAs is larger or equal to mPtrNumNonExitingCtas. - int32_t mMaxNumCtasInTokenDim{0}; - - // Either mBatchedM or mBatchedN must be set when mNumTokens == 0, otherwise not used. - // The number of tokens in each batch on the M dimension if batchM, - // otherwise not used. - // The number of elements in the array is mNumBatches. - // E.g. to implement a BMM with each batch having M tokens, one needs to set mBatchedM to - // {M, M, M, .. mNumBatches times ..} - std::vector mBatchedM{}; - // The number of tokens in each batch on the N dimension if batchN, - // otherwise not used. - // The number of elements in the array is mNumBatches. - // E.g. to implement a BMM with each batch having N tokens, one needs to set mBatchedN to - // {N, N, N, .. mNumBatches times ..} - std::vector mBatchedN{}; - - // The M dimension. - // It is the total number of tokens if A is the activation matrix. - // It is the total number of output channels if A is the weight matrix. - // ValidM/N/K by default assumes to be full range of M/N/K respectively. If we pad M/N/K due to - // alignment of other constraints, then we can specify ValidM/N/K to indicate the valid range. - int32_t mM{0}; - int32_t mValidM{0}; - // The N dimension. - // It is the total number of tokens if B is the activation matrix. - // It is the total number of output channels if B is the weight matrix. - int32_t mN{0}; - int32_t mValidN{0}; - // The K dimension. It is the hidden dimension of the input matrices. - int32_t mK{0}; - int32_t mValidK{0}; - // The rank id of the current device in the multi-gpu space. - int32_t mRank{0}; - // The number of devices in tensor-parallel group. - int32_t mWorldSize{1}; - }; - - struct InputBuffers { - // The matrix A. The data type is controlled by options.mDtypeA. - // - // Sparsity is only supported with batchN. - // Let S be the sparsity ratio (1 for dense, 2 for sparse). - // - // If (routeAct == true && batchM), the shape is [M, K] - // Elseif (batchStrideInTokens > 0) - // If batchM: - // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), K]. - // Logical strides are [K, 1]. - // - // If batchN: - // If layoutA is MatrixLayout::MajorK - // Logical shape is [B, divUpMul(M, tileM), K / S]. - // Logical strides are [divUpMul(M, tileM) * K / S, K / S, 1]. - // If layoutA is MatrixLayout::MajorMn (sparsity not supported) - // Logical shape is [B, K, divUpMul(M, tileM)]. - // Logical strides are [K * divUpMul(M, tileM), divUpMul(M, tileM), 1]. - // If layoutA is MatrixLayout::BlockMajorK - // Logical shape is [B, K / S / blockK, divUpMul(M, tileM), blockK]. - // Logical strides are [K / S * divUpMul(M, tileM), divUpMul(M, tileM) * blockK, blockK, - // 1]. where blockK is 128B. - // Else // batchStrideInTokens == 0 - // If batchM: (sparsity not supported) - // Logical shape is [M, K]. - // Logical strides are [K, 1]. - // - // If batchN: - // Logical shape is [B, divUpMul(M, tileM), K / S]. - // Logical strides are [divUpMul(M, tileM) * K / S, K / S, 1]. - void const* mPtrA{nullptr}; - - // The block scaling factors to dequantize A. - // - // If (routeAct == true && batchM), the shape is [M, K / 16] - // Else - // If DeepSeek FP8 recipe is used: - // If transposeMmaOutput is false, shape is [K / 128, M]. - // Otherwise, shape is [M / 128, K / 128]. - // The rightmost dimension is contiguous in memory. - // - // If DeepSeek FP8 recipe is not used, but for MxFp{4,8}, MxInt4 and NvFp4 formats: - // The layout of scaling factors for A is always R128c4 - // M must be a multiple of 128. - // K must be a multiple of 64. - // The "logical" shape is: [paddedM, K / P], where P is the scaling block size. - // The R128c4 layout is: [paddedM / 128, K / P / 4, 512]. - // The shape we use for TMA is: [paddedM / 128, K / P / 4, 2, 256]. - // Where paddedM is M if (routeAct == true && batchM), or - // sum(divUpMul(M[bi], tileM) for bi in B) if batchM, - // otherwise divUpMul(M, tileM) * B. - // Dtype is Dtype::Fp32 if DeepSeek FP8 recipe is used, otherwise Dtype is Dtype::E4m3 for - // NvFp4, Dtype::UE8m0 for MxFp{4,8} formats, Dtype::Bfloat16 for MxInt4. - // - // Otherwise should be set to nullptr. - void const* mPtrSfA{nullptr}; - - // The per-token scaling factors from scale A. - // - // This is used for either: - // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is - // Dtype::Float32 - // * When the routing scales are applied to the input activations (only when output is not - // transposed). The dtype is Dtype::Bfloat16 - // - // if (batchM (A is activations)): - // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B)] - // - // if (batchN (A is weights)): - // Logical shape is [B, divUpMul(M, tileM)] - // - void const* mPtrPerTokenSfA{nullptr}; - - // The matrix B. The data type is controlled by options.mDtypeB. - // - // If (routeAct == true && batchN), the shape is [N, K] - // - // Else if (batchStrideInTokens > 0) - // If batchN: - // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), K]. - // Logical strides are [K, 1]. - // - // If batchM: - // If layoutB is MatrixLayout::MajorK - // Logical shape is [B, divUpMul(N, tileN), K]. - // Logical strides are [divUpMul(N, tileN) * K, K, 1]. - // If layoutB is MatrixLayout::MajorMn - // Logical shape is [B, K, divUpMul(N, tileN)]. - // Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN), 1]. - // If layoutB is MatrixLayout::BlockMajorK - // Logical shape is [B, K / blockK, divUpMul(N, tileN), blockK]. - // Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN) * blockK, blockK, 1]. - // where blockK is 128B. - // - // Else // batchStrideInTokens == 0 - // If batchN: - // Logical shape is [N, K]. - // Logical strides are [K, 1]. - // - // If batchM: - // Logical shape is [B, divUpMul(N, tileN), K]. - // Logical strides are [divUpMul(N, tileN) * K, K, 1]. - void const* mPtrB{nullptr}; - - // The scaling factors to dequantize B. - // - // - // - // Else - // If DeepSeek FP8 recipe is used: - // If transposeMmaOutput is false, shape is [paddedN / 128, K / 128]. - // Otherwise, shape is [K / 128, paddedN]. - // The rightmost dimension is contiguous in memory. - // - // If DeepSeek FP8 recipe is not used, but for MxFp{4,8} and NvFp4 formats: - // If the layout is R128c4, - // paddedN must be a multiple of 128. - // K must be a multiple of 64. - // The R128c4 layout is: [paddedN / 128, K / P / 4, 512], where P is the scaling block - // size. The shape we use for TMA is: [paddedN / 128, K / P / 4, 2, 256] - // - // If the layout is R8c4, - // paddedN must be a multiple of 8. - // K must be a multiple of 64. - // The R8c4 layout is: [paddedN / 8, K / P / 4, 32], where P is the scaling block size. - // The shape we use for TMA is: [paddedN / 8, K / P / 4 / repeats, repeats * 32] - // where repeats = min(tileK / P / 4, 8) - // - // where paddedN is N if (routeAct == true && batchN), - // or sum(divUpMul(N[bi], tileN) for bi in B) if batchN, - // otherwise divUpMul(N, TileN) * B. - // - // Dtype is Dtype::Fp32 if DeepSeek FP8 recipe is used, otherwise Dtype::E4m3. - // - // Otherwise should be set to nullptr. - void const* mPtrSfB{nullptr}; - - // The per-token scaling factors from scale B. - // - // This is used for either: - // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is - // Dtype::Float32 - // * When the routing scales are applied to the input activations (only when output is - // transposed). The dtype is Dtype::Bfloat16 - // - // if (batchM (B is weights)): - // Logical shape is [B, divUpMul(N, tileN)] - // - // if (batchN (B is activations)): - // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B)] - void const* mPtrPerTokenSfB{nullptr}; - - // The sparsity information of A, if structured sparsity is used. - // Only supported for batchN (A is weights). - // - // When sparsityA is Any_2_4: - // 2 elements are non-zero in any chunk of 4 elements. - // A 4-bit index indicates the position of the non-zero elements. - // The shape in Uint8 is: [B, divUpMul(M, tileM), K / 8] - // (two 4-bit indices packed into one UInt8) - // - // When sparsityA is Pairwise_4_8: - // 4 elements are non-zero in any chunk of 8 elements. - // The zero and non-zero elements are grouped in pairs. - // A 4-bit index indicates the position of the non-zero pairs. - // The shape in Uint8 is: [B, divUpMul(M, tileM), K / 16] - // (two 4-bit indices packed into one UInt8) - // - // If sparsityA is Dense, this should be set to nullptr. - void const* mPtrSparsityInfoA{nullptr}; - - // The bias applied after the GEMM and before the activation function. - // The bias is applied before applying the global scaling factor. I.e. - // C = act(A * B + bias') * scaleC - // scaleC = dequantA * dequantB * quantC - // Thus, the bias' = bias / (dequantA * dequantB), where the bias is the original bias. - // - // If batchM, BiasType must be N, and bias shape is [B, N]. - // The bias is broadcasted along the M dimension. - // - // If batchN BiasType must be M, and bias shape is [B, M]. - // The bias is broadcasted along the N dimension. - // - // The dtype is float32. - void const* mPtrBias{nullptr}; - - // The output tensor scaling factor for Fp8 (not DeepSeek FP8) and NvFp4 quantization. - // TensorRT-LLM API requires a scaling factor on the device. - // scaleC = dequantA * dequantB * quantC, - // where dequantA is global dequantization scaling factor of A - // if dtypeA is FP8, it transforms the range from [-448, 448] to [-amaxA, amaxA] - // if dtypeA is NvFp4, it transforms the range from [-448 * 6, 448 * 6] to [-amaxA, amaxA], - // otherwise it is 1. - // dequantB is defined similarly to dequantA. - // quantC is the quantization scaling factor of C. - // if dtypeC is FP8, it transforms the range from [-amaxC, amaxC] to [-448, 448] - // if dtypeC is NvFp4, it transforms the range from [-amaxC, amaxC] to [-448 * 6, 448 * 6], - // otherwise it is 1. - // Shape is [B]. - float const* mPtrScaleC{nullptr}; - - // The pre-activation scaling factor (typically dequantA * dequantB) for non-gated non-linear - // activation. - // Only used when non-linear activation is applied (e.g., GELU, Relu2). - // When used, scaleC should be quantScaleC only, and this scale is applied before the - // activation. Shape is [B]. - float const* mPtrScaleAct{nullptr}; - - // The output gate scale for Fp8 (not DeepSeek FP8) and NvFp4 quantization. - // TensorRT-LLM API requires a scaling factor on the device. - // scaleGate = dequantA * dequantB, - // where dequantA is global dequantization scaling factor of A - // if dtypeA is FP8, it transforms the range from [-448, 448] to [-amaxA, amaxA] - // if dtypeA is NvFp4, it transforms the range from [-448 * 6, 448 * 6] to [-amaxA, amaxA], - // otherwise it is 1. - // dequantB is defined similarly to dequantA. - // Shape is [B]. - float const* mPtrScaleGate{nullptr}; - - // The clamp limit for the accumulator before applying the activation. - // Shape is [B]. - // Clamp is INF if nullptr. - // When the input is FP8 or NVFP4, the clamp has to be scaled by limit' = limit / dequantAb. - // If applied on SwiGlu, it will be: - // - // x_glu = x_glu.clamp(min=None, max=limit) - // x_linear = x_linear.clamp(min=-limit, max=limit) - // - // The given clamp limit applies to the dequantized values, so the order of operations would - // look something like this: - // - // x0 = x0 * dqAb - // x0 = clamp(x0, none, limit) - // x0 = x0 * sigmoid(alpha * x0) - // x1 = dqAb * x1 - // x1 = clamp(x1, -limit, limit) - // out = qC * (x1 + beta) * x0 - // - // Given that the dqAb and qC are combined into scaleC, we can bring the dqAb into the clamp - // limit and apply the clamping prior to dequantization: - // - // x0 = clamp(x0, none, limit / dqAb) - // x0 = x0 * dqAb - // x0 = x0 * sigmoid(alpha * x0) - // x1 = clamp(x1, -limit / dqAb, limit / dqAb) - // scaleC = dqAb * qC - // beta' = beta / dqAb - // out = scaleC * (x1 + beta') * x0 - // - // Note this assumes that dequantScaleAb == scaleGate which is true in TRT-LLM MoE use-case - // - float const* mPtrClampLimit{nullptr}; - - // The alpha and beta for SwiGlu or GeGlu. - // gatedActivation <- (x0 + beta) * activation(x1, alpha) - // Shape is [B]. - // Alpha is 1.f if nullptr. - // Beta is 0.f if nullptr. - // The formula for SwiGlu (for GeGlu, replace sigmoid with phi): - // - // out_glu = x_glu * torch.sigmoid(alpha * x_glu) * (x_linear + beta) - // - // The beta is added before applying the global scaling factor. I.e. - // x_linear = (x_linear + beta') * scaleC - // Thus, the beta' = beta / (dequantA * dequantB), where the beta is the original beta. - float const* mPtrGatedActAlpha{nullptr}; - float const* mPtrGatedActBeta{nullptr}; - - // Param is used when the kernel is configured with -routeAct true. - // The inputs are not padded, but the outputs are padded to divUpMul(M[bi], tileM) for batchM or - // divUpMul(N[bi], tileN) for batchN. - // If -routeAct is false, the params are not used and should be set to zero. - - // The routeMap for the input tokens. - // Map of expanded token index (counting the previous padded tokens) to the batch index - // the token belongs to. - // The shape is - // [divUpMul(numTokens + numBatches * (tileM/N - 1), tileM/N)] - // The dtype is int32_t. - // - // There are 3 tokens [0, 1, 2] such that [0, 1] belong to batch [B0] and [2] to batch [B1]. - // Let's assume that the padded size is 4. - // - // The expanded indices for tokens [0, 1, 2] are: - // expandedIdx[0] = 0 - // expandedIdx[1] = 1 - // expandedIdx[2] = divUpMul(2, 4) + 0 = 4 - // - // The route map is [B0, B0, X, X, B1, X, X, X] where X could be any value. - int32_t const* mPtrRouteMap; - - ////////////////////////////////////////////////////////////////////////////////////////////////// - // - // Batching information parameters. - // - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // In some cases, some CTAs must early-exit. E.g. when the grid size is set statically, but the - // actual workload is decided at runtime. This element on the device contains the number of CTAs - // that do not early-exit. The number corresponds to the X dim of the grid when the output is - // not transposed (i.e. batchM). To the Y dim, otherwise. The size is 1 and the dtype is - // int32_t. Used if isStaticBatch == false, otherwise set to nullptr. The pointer points to a - // scalar and the dtype is int32_t. The pointed value must be >= 0. - int32_t const* mPtrNumNonExitingCtas; - - // Pointer to total number of padded tokens. - // Computed as - // int32_t totalNumPaddedTokens{0}; - // for (int bi = 0; bi < options.mNumBatches; bi++) { - // totalNumPaddedTokens += batchM ? divUpMul(options.mBatchedM[bi], options.mTileM * - // options.mClusterDimX) - // : divUpMul(options.mBatchedN[bi], options.mTileN); - // } - // The size is 1 and the dtype is int32_t. - // If isStaticBatch == true, ptrTotalNumPaddedTokens should be set to nullptr and - // totalNumPaddedTokens is used. - int32_t const* mPtrTotalNumPaddedTokens; - - // Pointer to the map from the CTA index (in X/Y dim) to the batch index. - // Maps CTA index in batch dim (i.e. blockDim.x if batchM, otherwise blockDim.y) - // to batch index. - // E.g. with listM = 128,255,32 and tileM = 128, should be equal to - // ctaIdxXyToBatchIdx = [0, 1, 1, 2] - // If isStaticBatch == true, ptrCtaIdxXyToBatchIdx should be set to nullptr and - // ctaIdxXyToBatchIdx is used. - // The shape is - // [divUp(numTokens + numBatches * (tileM/N - 1), tileM/N)] - int32_t const* mPtrCtaIdxXyToBatchIdx; - - // Pointer from the CTA index X/Y to the expanded tile index where the expanded tile index is - // computed as: - // - // int expandedIdx = 0; - // for (int bi = 0; bi < batchIdx-1; ++bi) { - // expandIdx = divUpMul(numTokens[bi], TileM/N); - // } - // expandIdx += - // E.g. with numTokens = [128,255,32] and tileM = 128, should be equal to - // ptrCtaIdxXyToMnLimit = [128, 256, 383, 416] - // The shape is - // [divUp(numTokens + numBatches * (tileM/N - 1), tileM/N)] - int32_t const* mPtrCtaIdxXyToMnLimit; - }; - - struct OutputBuffers { - // The output matrix C. The data type is controlled by options.mDtypeC. - // - // If batchM: - // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), N]. - // Logical strides are [N, 1]. - // - // If batchN: - // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), M]. - // Logical strides are [M, 1]. - void* mPtrC{nullptr}; - - // The output block scaling factors for C. - // - // If MxFp{4,8} and NvFp4 formats are used, - // The "logical" shape is (P is the scaling block size): - // if batchM: [paddedM, N / P] - // if batchN: [paddedN, M / P] - // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B), - // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). - // - // If the layout is R128c4, - // paddedOuter must be a multiple of 128. - // inner must be a multiple of 64. - // The R128c4 layout is: [paddedOuter / 128, inner / P / 4, 512] - // The shape we use for TMA is: [paddedOuter / 128, inner / P / 4, 2, 256] - // where inner = N if batchM, otherwise M. - // where paddedOuter = paddedM if batchM, otherwise paddedN. - // - // If the layout is R8c4, - // paddedOuter must be a multiple of 8. - // inner must be a multiple of 64. - // The R8c4 layout is: [paddedOuter / 8, inner / P / 4, 32] - // The shape we use for TMA is: [paddedOuter / 8, inner / P / 4 / repeats, repeats * 32] - // where repeats = min(tileInner / P / 4, 8), - // where tileInner = tileN if batchM, otherwise tileM, - // where paddedOuter = paddedM if batchM, otherwise paddedN. - // where inner = N if batchM, otherwise M. - // - // The dtype is Dtype::E4m3. - // - // If DeepSeek FP8 quantization recipe is used, - // If batchM: - // The shape is [N / 128, paddedM], - // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B). - // If batchN: - // The shape is [M / 128, paddedN], - // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). - // The rightmost dimension is contiguous in memory. - // The dtype is Dtype::Float32. - void* mPtrSfC{nullptr}; - }; - - ProblemDimensions mProblemDimensions; - InputBuffers mInputBuffers; - OutputBuffers mOutputBuffers; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// BatchedGemmInterface -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -class BatchedGemmInterface { - public: - using ModuleCache = std::unordered_map>; - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - BatchedGemmInterface(bool const exportsCubin = false, int32_t const numRotations = 1) - : mExportsCubin(exportsCubin), mNumRotations(numRotations) {} - - ////////////////////////////////////////////////////////////////////////////////////////////////// - -#ifndef TLLM_GEN_EXPORT_INTERFACE - // Generates and compiles the kernel using either nvcc or nvrtc. - BatchedGemmConfig generateAndCompileKernel(BatchedGemmConfig const& batchedGemmConfig) const; -#endif - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Launch the cubin from the provided config. It calls all necessary memsets for internal buffers. - // Provided config must be validated with isValidConfig before the call. - int32_t run(BatchedGemmConfig const& config, void* workspace, - BatchedGemmData const& batchedGemmData, void* cudaStream, - int32_t /*multiProcessorCount*/, bool usePdl = true, - std::optional> moduleCache = std::nullopt) { - // Get options from config and data. - auto options = getOptionsFromConfigAndData(config, batchedGemmData); - - bool const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM; - bool const useDeepSeekFp8 = options.mUseDeepSeekFp8 && options.mDtypeA == tg::Dtype::E4m3 && - options.mDtypeB == tg::Dtype::E4m3; - - auto workspaceSizes = getWorkspaceSizesInBytes(config, batchedGemmData); - float* dPtrRowMax{nullptr}; - uint32_t* dPtrRowMaxBars{nullptr}; - - // Set the completion barriers to 0 if needed. - if (useDeepSeekFp8 && options.mFusedAct) { - dPtrRowMax = reinterpret_cast(alignPtr(reinterpret_cast(workspace), 1024)); - dPtrRowMaxBars = reinterpret_cast( - alignPtr(reinterpret_cast(dPtrRowMax) + workspaceSizes[0], 1024)); - auto err = cudaMemsetAsync((void*)dPtrRowMaxBars, 0x00, workspaceSizes[1], - reinterpret_cast(cudaStream)); - if (err != cudaSuccess) { - return 1; - } - } - - auto [numCtaBatch, numCtaTile, numCtaInner] = - getGridDim(options, batchedGemmData.mProblemDimensions.mMaxNumCtasInTokenDim); - - auto kernelParams = KernelParamsSetup::setKernelParams( - options, batchM, batchedGemmData.mInputBuffers.mPtrA, batchedGemmData.mInputBuffers.mPtrB, - batchedGemmData.mOutputBuffers.mPtrC, batchedGemmData.mInputBuffers.mPtrSfA, - batchedGemmData.mInputBuffers.mPtrSfB, batchedGemmData.mInputBuffers.mPtrPerTokenSfA, - batchedGemmData.mInputBuffers.mPtrPerTokenSfB, - batchedGemmData.mInputBuffers.mPtrSparsityInfoA, batchedGemmData.mInputBuffers.mPtrBias, - batchedGemmData.mOutputBuffers.mPtrSfC, batchedGemmData.mInputBuffers.mPtrScaleC, - batchedGemmData.mInputBuffers.mPtrScaleAct, batchedGemmData.mInputBuffers.mPtrScaleGate, - batchedGemmData.mInputBuffers.mPtrClampLimit, - batchedGemmData.mInputBuffers.mPtrGatedActAlpha, - batchedGemmData.mInputBuffers.mPtrGatedActBeta, batchedGemmData.mInputBuffers.mPtrRouteMap, - dPtrRowMax, dPtrRowMaxBars, batchedGemmData.mInputBuffers.mPtrNumNonExitingCtas, - batchedGemmData.mInputBuffers.mPtrTotalNumPaddedTokens, - batchedGemmData.mInputBuffers.mPtrCtaIdxXyToBatchIdx, - batchedGemmData.mInputBuffers.mPtrCtaIdxXyToMnLimit, numCtaBatch); - - // The size of the grid. - auto grid = getLaunchGrid(options, batchedGemmData.mProblemDimensions.mMaxNumCtasInTokenDim); - - BatchedGemmConfig batchedGemmConfig = config; -#ifndef TLLM_GEN_EXPORT_INTERFACE - // Generate and compile the kernel if data is not provided. - if (config.mData == nullptr) { - batchedGemmConfig = generateAndCompileKernel(batchedGemmConfig); - } - TLLM_CHECK_ERROR(batchedGemmConfig.mCudaRunner != nullptr, "CudaRunner is not set"); - batchedGemmConfig.mCudaRunner->run((void*)&kernelParams, (void*)cudaStream, grid, - /* cluster */ {}, - /* instanceId */ batchedGemmConfig.mInstanceIdx); - return 0; -#endif - - CUmodule cuModule; - CUfunction cuFunction; - - if (moduleCache.has_value()) { - ModuleCache& moduleCacheRef = moduleCache.value().get(); - - // Modules are associated with a specific context, so the context is included in the key - CUcontext ctx; - unsigned long long ctxId; - cuCtxGetCurrent(&ctx); - cuCtxGetId(ctx, &ctxId); - - // Reinterpret the ctxId as a string to avoid needing a custom hash or converting it to a - // string in decimal representation. - std::string const ctxName = - std::string(reinterpret_cast(&ctxId), sizeof(unsigned long long) / sizeof(char)); - std::string const funcName = std::string(batchedGemmConfig.mFunctionName); - auto const moduleKey = ctxName + funcName; - auto module = moduleCacheRef.find(moduleKey); - - // Use cache if module is found, otherwise load and insert into cache - if (module != moduleCacheRef.end()) { - cuFunction = std::get<1>(module->second); - } else { - gemm::loadCubinData(&cuModule, batchedGemmConfig); - cuModuleGetFunction(&cuFunction, cuModule, batchedGemmConfig.mFunctionName); - moduleCacheRef.insert(std::make_pair(moduleKey, std::make_tuple(cuModule, cuFunction))); - } - } else { - gemm::loadCubinData(&cuModule, batchedGemmConfig); - cuModuleGetFunction(&cuFunction, cuModule, batchedGemmConfig.mFunctionName); - } - - // Prepare the grid/block. - dim3 block3{static_cast(batchedGemmConfig.mNumThreadsPerCTA), - static_cast(1), static_cast(1)}; - dim3 grid3{(grid.size() > 0 ? static_cast(grid[0]) : 1u), - (grid.size() > 1 ? static_cast(grid[1]) : 1u), - (grid.size() > 2 ? static_cast(grid[2]) : 1u)}; - // Prepare the cluster size. - dim3 cluster3{static_cast(options.mClusterDimX), - static_cast(options.mClusterDimY), - static_cast(options.mClusterDimZ)}; - - // Whether PDL can safely be enabled - const bool pdlSafe = batchedGemmConfig.mOptions.mGridWaitForPrimaryRouting || - batchedGemmConfig.mOptions.mGridWaitForPrimaryEarlyExit || - batchedGemmConfig.mOptions.mGridWaitForPrimaryA || - batchedGemmConfig.mOptions.mGridWaitForPrimaryB; - - // Run the kernel. - auto result = trtllm::gen::launchKernel((void*)&kernelParams, cudaStream, - batchedGemmConfig.mSharedMemSize, cuFunction, block3, - grid3, cluster3, usePdl && pdlSafe); - if (result != CUDA_SUCCESS) { - return result; - } - // If a module cache has not been given, unload the module to avoid leaking - if (!moduleCache.has_value()) { - cuModuleUnload(cuModule); - } - return 0; - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Initializes the buffers before the world sync. Must be called before run. - int32_t runInitBeforeWorldSync(BatchedGemmConfig const& /* config */, - BatchedGemmData const& /* data */, void* /* cudaStream */) const { - return 0; - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - size_t getWorkspaceSizeInBytes(BatchedGemmConfig const& config, - BatchedGemmData const& data) const { - auto workspaceSizes = getWorkspaceSizesInBytes(config, data); - auto size = std::accumulate(workspaceSizes.begin(), workspaceSizes.end(), 0); - // Additional 1023 bytes to align the pointer to 1024 - return size > 0 ? size + 1023 : 0; - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Returns the list of all available cubin configurations - BatchedGemmConfig const* getBatchedGemmConfigs() const { -#ifdef TLLM_GEN_EXPORT_INTERFACE - return tensorrt_llm::kernels::tllmGenBatchedGemmList; -#else - return nullptr; -#endif - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Returns the number of available cubin configurations - size_t getNumBatchedGemmConfigs() const { -#ifdef TLLM_GEN_EXPORT_INTERFACE - return tensorrt_llm::kernels::tllmGenBatchedGemmListLen; -#else - return 0; -#endif - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Returns the grid dimensions of the current kernel. - std::tuple getGridDim( - BatchedGemmOptions const& options, - std::optional maxNumCtasInBatchDim = std::nullopt) const { - bool const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM; - - int32_t numCtasBatch{0}; - // For normal BMM, mNumTokens == 0 and the number of CTAs is known to host. - if (options.mIsStaticBatch) { - for (int32_t bi = 0; bi < options.mNumBatches; ++bi) { - numCtasBatch += - batchM ? gemm::divUp(options.mBatchedM[bi], options.mTileM * options.mClusterDimX) * - options.mClusterDimX - : gemm::divUp(options.mBatchedN[bi], options.mTileN); - } - } - // For MoE, mNumTokens != 0 and the number of CTAs is known only at runtime. - // We launch maximally possible number of CTAs and use ptrNumNonExitingCtas to determine the - // actual number of CTAs to run. - else if ((options.mEnablesEarlyExit || options.mEnablesDelayedEarlyExit) && - options.mNumTokens != 0) { - assert(maxNumCtasInBatchDim.has_value() && - "maxNumCtasInBatchDim must be provided when options.mNumTokens != 0"); - numCtasBatch = maxNumCtasInBatchDim.value(); - } else { - throw std::invalid_argument("Invalid combination of options"); - } - - if (batchM) { - numCtasBatch = gemm::divUpMul(numCtasBatch, options.mClusterDimX); - } else { - numCtasBatch = gemm::divUpMul(numCtasBatch, options.mClusterDimY); - } - - int32_t numCtasTile = - batchM ? gemm::divUp(options.mN, options.mTileN) : gemm::divUp(options.mM, options.mTileM); - if (batchM) { - numCtasTile = gemm::divUpMul(numCtasTile, options.mClusterDimY); - } else { - numCtasTile = gemm::divUpMul(numCtasTile, options.mClusterDimX); - } - int32_t const numCtasInner = options.mNumSlicesForSplitK; - return std::make_tuple(numCtasBatch, numCtasTile, numCtasInner); - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Returns the number of CTAs of the current kernel. - std::vector getLaunchGrid( - BatchedGemmOptions const& options, - std::optional maxNumCtasInBatchDim = std::nullopt) const { - auto [numCtaBatch, numCtaTile, numCtaInner] = getGridDim(options, maxNumCtasInBatchDim); - bool const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM; - std::vector grid = batchM ? std::vector{numCtaBatch, numCtaTile, numCtaInner} - : std::vector{numCtaTile, numCtaBatch, numCtaInner}; - return grid; - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Returns the number of CTAs of the current kernel. - int32_t getNumCtas(BatchedGemmOptions const& options, - std::optional maxNumCtasInBatchDim = std::nullopt) const { - auto grid = getLaunchGrid(options, maxNumCtasInBatchDim); - return grid[0] * grid[1] * grid[2]; - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Creates GemmOptions from kernel and data. - BatchedGemmOptions getOptionsFromConfigAndData(BatchedGemmConfig const& config, - BatchedGemmData const& data) const { - BatchedGemmOptions options; - options = config.mOptions; - options.mM = data.mProblemDimensions.mM; - options.mN = data.mProblemDimensions.mN; - options.mK = data.mProblemDimensions.mK; - options.mValidM = data.mProblemDimensions.mValidM; - options.mValidN = data.mProblemDimensions.mValidN; - options.mValidK = data.mProblemDimensions.mValidK; - options.mBatchedM = data.mProblemDimensions.mBatchedM; - options.mBatchedN = data.mProblemDimensions.mBatchedN; - options.mBatchMode = data.mProblemDimensions.mBatchM ? BatchedGemmOptions::BatchMode::BatchM - : BatchedGemmOptions::BatchMode::BatchN; - options.mNumBatches = data.mProblemDimensions.mNumBatches; - options.mNumTokens = data.mProblemDimensions.mNumTokens; - return options; - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Returns true if the configuration of the cubin can be executed for the given params. - bool isValidConfig(BatchedGemmConfig const& config, BatchedGemmData const& data) const { - // Get options from config and data. - auto options = getOptionsFromConfigAndData(config, data); - - // Check options without modifications. - return checkAndUpdateBatchedGemmOptions(options, config.mSm, - /* updateOptions */ false); - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - private: - ////////////////////////////////////////////////////////////////////////////////////////////////// - - template - inline Dtype* alignPtr(Dtype* ptr, int64_t alignment) const { - assert((alignment & (alignment - 1)) == 0 && "Alignment must be a power of 2"); - return reinterpret_cast((reinterpret_cast(ptr) + alignment - 1) & - ~(alignment - 1)); - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Returns the size of the workspace buffers in bytes - std::vector getWorkspaceSizesInBytes(BatchedGemmConfig const& config, - BatchedGemmData const& data) const { - std::vector workspaceSizes; - - // Get options from config and data. - auto options = getOptionsFromConfigAndData(config, data); - - if (options.mUseDeepSeekFp8 && options.mFusedAct) { - int32_t totalNumPaddedTokens = 0; - auto const batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM; - if (!options.mEnablesEarlyExit || options.mNumTokens == 0) { - for (int32_t bi = 0; bi < options.mNumBatches; ++bi) { - totalNumPaddedTokens += - batchM ? gemm::divUpMul(options.mBatchedM[bi], options.mTileM * options.mClusterDimX) - : gemm::divUpMul(options.mBatchedN[bi], options.mTileN); - } - } else { - // Get tile in token dim. - auto tileTokensDim = batchM ? options.mTileM * options.mClusterDimX : options.mTileN; - totalNumPaddedTokens = data.mProblemDimensions.mMaxNumCtasInTokenDim * tileTokensDim; - } - - // Get options from config. - auto& options = config.mOptions; - - int const tokenTile = batchM ? options.mTileM * options.mClusterDimX : options.mTileN; - - auto const numTokens = totalNumPaddedTokens; - auto const intermediateDim = batchM ? options.mN : options.mM; - auto const intermediateTile = batchM ? options.mTileN : options.mTileM; - - auto const numBytesRowMax = intermediateDim * totalNumPaddedTokens / 128 * sizeof(float); - - auto const numTilesToken = numTokens / tokenTile; - auto const numTilesInt = intermediateDim / intermediateTile; - auto const numBytesRowMaxBars = numTilesToken * numTilesInt / 2 * sizeof(uint32_t); - - // TODO: do we need to pad to 1024? - workspaceSizes.push_back(getSizePaddedToAlignment(numBytesRowMax, 1024)); - workspaceSizes.push_back(getSizePaddedToAlignment(numBytesRowMaxBars, 1024)); - } - - return workspaceSizes; - } - - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Returns the size padded to the alignment - size_t getSizePaddedToAlignment(size_t size, size_t alignment) const { - assert((alignment & (alignment - 1)) == 0); - return (size + alignment - 1) & ~(alignment - 1); - } - ////////////////////////////////////////////////////////////////////////////////////////////////// - - private: - // Whether to export the cubin file. - bool mExportsCubin; - // The number of rotations. - int32_t mNumRotations; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace batchedGemm - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h deleted file mode 100644 index 9c571d0258..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/BatchedGemmOptions.h +++ /dev/null @@ -1,525 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include - -#include "BatchedGemmEnums.h" -#include "GemmGatedActOptions.h" -#include "GemmOptions.h" - -#ifndef TLLM_GEN_EXPORT_INTERFACE -#include "trtllm/gen/CudaRunner.h" -#include "trtllm/gen/GenCtx.h" -#else -#include - -#define TLLM_CHECK_ERROR(cond, ...) \ - if (!(cond)) { \ - printArgs(__VA_ARGS__); \ - printArgs("\n"); \ - return false; \ - } - -#define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) - -#define TLLM_CHECK_ERROR_FMT(cond, ...) TLLM_CHECK_ERROR(cond, __VA_ARGS__) - -#define TLLM_CHECK_WARNING(cond, ...) \ - if (!(cond)) { \ - printArgs(__VA_ARGS__); \ - printArgs("\n"); \ - return false; \ - } - -#define TLLM_LOG_WARNING(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) - -#define TLLM_LOG_INFO(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) - -#endif // TLLM_GEN_EXPORT_INTERFACE - -namespace batchedGemm { - -namespace trtllm { -namespace gen { -class CudaRunner; -class GenCfg; -} // namespace gen -} // namespace trtllm - -namespace batchedGemm { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace tg = trtllm::gen; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// We do not differentiate between BatchedGemmOptions and BatchedGemmGatedActOptions for simplicity. -// We inherit from GemmGatedActOptions, which is inherited from -// GemmOptions to get GemmOptions and GemmGatedActOptions at the same time. -struct BatchedGemmOptions : public gemmGatedAct::GemmGatedActOptions { - // Dtor. Allow down-casting. - virtual ~BatchedGemmOptions() = default; - - enum class BatchMode { BatchM, BatchN }; - - BatchedGemmOptions() = default; - // FIXME We create explicit constructor with all options to WAR stubgen issue in TRT-LLM. - BatchedGemmOptions( - gemm::AllReduceAlgo allReduceAlgo, gemm::BiasType biasType, int blockK, int clusterDimX, - int clusterDimY, int clusterDimZ, gemm::CtaSwizzleType ctaSwizzleType, tg::Dtype dtypeAcc, - tg::Dtype dtypeA, tg::Dtype dtypeB, tg::Dtype dtypeC, tg::Dtype dtypeMmaA, - tg::Dtype dtypeMmaB, gemm::EltwiseActType eltwiseActType, bool enablesEarlyExit, - bool enablesDelayedEarlyExit, bool enablesGlobalPtxKnobs, int epilogueLdtmDps, - int epilogueLdtmBits, int epilogueTileM, int epilogueTileN, bool fuseUtccpWithUtcmma, - bool gridTriggerSecondaryA, bool gridTriggerSecondaryB, bool gridWaitForPrimaryEarlyExit, - bool gridWaitForPrimaryA, bool gridWaitForPrimaryB, bool hoistLoadTaskInit, - bool hoistMmaTaskTryWaits, int k, gemm::KernelTraits kernelTraits, gemm::MatrixLayout layoutA, - gemm::MatrixLayout layoutB, int m, int mmaK, tg::MmaKind mmaKind, int mmaM, int mmaN, - bool mockAllReduce, int n, int numEpilogueWarps, int numRegsCastAWarps, - int numRegsCopySfLdsSttm, int numRegsCopySparsityInfo, int numRegsPerThreadEpilogueWarp, - int numRegsPerThreadNonEpilogueWarp, int numSlicesForSplitK, int numSlicesForSliceK, - int numStages, int numStagesMma, int numStagesMmaWithinWorkTile, - int numStagesMmaAcrossWorkTile, int numStagesWorkId, bool outputDebugTensors, bool patchF2fp, - int32_t sfBlockSizeA, int32_t sfBlockSizeB, int32_t sfBlockSizeC, tg::SfLayout sfLayoutA, - tg::SfLayout sfLayoutB, tg::SfLayout sfLayoutC, int32_t sfReshapeFactor, bool sliceK, - tg::Sparsity sparsityA, gemm::SplitK splitK, int tileK, int tileM, int tileN, - gemm::TileScheduler tileScheduler, bool transposeMmaOutput, bool useCustomMmaSchedule, - bool useDeepSeekFp8, bool useHoistTryWaitForCustomMmaSchedule, bool useMaxTmemOverlap, - bool usePerTokenSfA, bool usePerTokenSfB, bool useShuffledMatrix, bool useTmaStore, - bool useTwoTmaLoadWarps, bool useTwoMmaWarps, bool useUnrollLoop2xForMma, int validM, - int validN, int validK, int worldSize, - // GemmGatedActOptions - gemmGatedAct::ActType actType, bool clampBeforeAct, - // BatchedGemmOptions - std::vector batchedM, std::vector batchedN, BatchMode batchMode, - int32_t batchStrideInTokens, bool fusedAct, bool gridWaitForPrimaryRouting, - bool isStaticBatch, bool isUniformNumTokensPerBatch, int numBatches, - int numRegsPerThreadLoadA, int numRegsPerThreadLoadB, int numRegsPerThreadLoadSfA, - int numRegsPerThreadLoadSfB, int numTokens, int numWarpsLoadA, int numWarpsLoadB, - int numWarpsLoadSfA, int numWarpsLoadSfB, RouteImpl routeImpl, - std::optional routeSfsImpl, bool useTmaOobOpt) - : gemmGatedAct::GemmGatedActOptions( - gemm::GemmOptions( - allReduceAlgo, biasType, blockK, clusterDimX, clusterDimY, clusterDimZ, - ctaSwizzleType, dtypeAcc, dtypeA, dtypeB, dtypeC, dtypeMmaA, dtypeMmaB, - eltwiseActType, enablesEarlyExit, enablesDelayedEarlyExit, enablesGlobalPtxKnobs, - epilogueLdtmDps, epilogueLdtmBits, epilogueTileM, epilogueTileN, - fuseUtccpWithUtcmma, gridTriggerSecondaryA, gridTriggerSecondaryB, - gridWaitForPrimaryEarlyExit, gridWaitForPrimaryA, gridWaitForPrimaryB, - hoistLoadTaskInit, hoistMmaTaskTryWaits, k, kernelTraits, layoutA, layoutB, m, mmaK, - mmaKind, mmaM, mmaN, mockAllReduce, n, numEpilogueWarps, numRegsCastAWarps, - numRegsCopySfLdsSttm, numRegsCopySparsityInfo, numRegsPerThreadEpilogueWarp, - numRegsPerThreadNonEpilogueWarp, numSlicesForSplitK, numSlicesForSliceK, numStages, - numStagesMma, numStagesMmaWithinWorkTile, numStagesMmaAcrossWorkTile, - numStagesWorkId, outputDebugTensors, patchF2fp, sfBlockSizeA, sfBlockSizeB, - sfBlockSizeC, sfLayoutA, sfLayoutB, sfLayoutC, sfReshapeFactor, sliceK, sparsityA, - splitK, tileK, tileM, tileN, tileScheduler, transposeMmaOutput, - useCustomMmaSchedule, useDeepSeekFp8, useHoistTryWaitForCustomMmaSchedule, - useMaxTmemOverlap, usePerTokenSfA, usePerTokenSfB, useShuffledMatrix, useTmaStore, - useTwoTmaLoadWarps, useTwoMmaWarps, useUnrollLoop2xForMma, validM, validN, validK, - worldSize), - actType, clampBeforeAct), - mBatchedM(batchedM), - mBatchedN(batchedN), - mBatchMode(BatchMode(batchMode)), - mBatchStrideInTokens(batchStrideInTokens), - mFusedAct(fusedAct), - mGridWaitForPrimaryRouting(gridWaitForPrimaryRouting), - mIsStaticBatch(isStaticBatch), - mIsUniformNumTokensPerBatch(isUniformNumTokensPerBatch), - mNumBatches(numBatches), - mNumRegsPerThreadLoadA{numRegsPerThreadLoadA}, - mNumRegsPerThreadLoadB{numRegsPerThreadLoadB}, - mNumRegsPerThreadLoadSfA{numRegsPerThreadLoadSfA}, - mNumRegsPerThreadLoadSfB{numRegsPerThreadLoadSfB}, - mNumTokens(numTokens), - mNumWarpsLoadA{numWarpsLoadA}, - mNumWarpsLoadB{numWarpsLoadB}, - mNumWarpsLoadSfA{numWarpsLoadSfA}, - mNumWarpsLoadSfB{numWarpsLoadSfB}, - mRouteImpl(routeImpl), - mRouteSfsImpl(routeSfsImpl), - mUseTmaOobOpt(useTmaOobOpt) {} - - // Batched M-dimensions of GEMM. - std::vector mBatchedM; - // Batched N-dimensions of GEMM. - std::vector mBatchedN; - // Whether batching M or N. - BatchMode mBatchMode{BatchMode::BatchM}; - // Stride between batches in tokens dimension for input matrix. - int32_t mBatchStrideInTokens{-1}; - // Whether to perform a fused gated activation. - bool mFusedAct{false}; - // Whether the loads that load from ptrRouteMap, ptrTotalNumPaddedTokens, - // ptrCtaIdxXyToBatchIdx, etc.. should wait on a grid dependency. - bool mGridWaitForPrimaryRouting{true}; - // Whether the batch size is static (i.e. known at kernel launch time). - bool mIsStaticBatch{true}; - // Whether the number of tokens in each entry of the batch is the same. - bool mIsUniformNumTokensPerBatch{false}; - // Number of Gemm batches. - int mNumBatches; - // Number of registers per thread for load A - int mNumRegsPerThreadLoadA{0}; - // Number of registers per thread for load B - int mNumRegsPerThreadLoadB{0}; - // Number of registers per thread for load SfA - int mNumRegsPerThreadLoadSfA{0}; - // Number of registers per thread for load SfB - int mNumRegsPerThreadLoadSfB{0}; - // Total number of tokens. - int mNumTokens{32}; - // Number of warps for load A - int mNumWarpsLoadA{0}; - // Number of warps for load B - int mNumWarpsLoadB{0}; - // Number of warps for load SfA - int mNumWarpsLoadSfA{0}; - // Number of warps for load SfB - int mNumWarpsLoadSfB{0}; - // Whether load the input tokens and do routing. - RouteImpl mRouteImpl{RouteImpl::NoRoute}; - // Routing logic for scaling factors. If not specified, mRouteImpl is used. - std::optional mRouteSfsImpl{std::nullopt}; - // Whether to use TMA out-of-bounds optimization to reduce wasted traffic. See details in - // BatchedGemm/KernelParamsDecl.h. - bool mUseTmaOobOpt{false}; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Check if the options are valid or not. -inline bool checkAndUpdateBatchedGemmOptions(BatchedGemmOptions& options, tg::CudaArch cudaArch, - bool updateOptions = true) { - bool isValid = true; - if (options.mUseTmaOobOpt && !options.mUseTwoTmaLoadWarps) { - if (updateOptions) { - // Since any routing (mRouteAct != NoRoute) requires mUseTwoTmaLoadWarps == true. - // Single TMA load warp is not the target use case for OOB optimization. - options.mUseTmaOobOpt = false; - } else if (!options.mUseTwoTmaLoadWarps) { - TLLM_CHECK_ERROR(false, "TMA OOB optimization requires two TMA load warps."); - return false; - } - } - if (options.mFusedAct) { - // ensure that we check the fused options as well - isValid = gemmGatedAct::checkAndUpdateGemmGatedActOptions(options, cudaArch, updateOptions); - } else { - isValid = gemm::checkAndUpdateGemmOptions(options, cudaArch, 1 /* tpGrpSize */, updateOptions); - } - - bool batchM = options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM; - if (updateOptions) { - if (batchM) { - if (options.mBatchedM.empty()) { - options.mBatchedM.push_back(128); - options.mBatchedM.push_back(256); - } - options.mNumBatches = options.mBatchedM.size(); - } else { - if (options.mBatchedN.empty()) { - options.mBatchedN.push_back(128); - options.mBatchedN.push_back(256); - } - options.mNumBatches = options.mBatchedN.size(); - } - } - - if (batchM) { - TLLM_CHECK_ERROR(!tg::isSparse(options.mSparsityA), "Sparsity is not supported with batchM."); - TLLM_CHECK_ERROR(options.mN > 0 && options.mK > 0, "N and K must be larger than 0"); - TLLM_CHECK_ERROR(options.mN >= options.mTileN, "N must be equal or larger than TileN."); - TLLM_CHECK_ERROR(options.mN % options.mTileN == 0, "N must be divisible by TileN."); - TLLM_CHECK_ERROR(!options.mTransposeMmaOutput, - "When batchM the MMA output has to be in row-major."); - } else { - TLLM_CHECK_ERROR(options.mM > 0 && options.mK > 0, "M and K must be larger than 0"); - TLLM_CHECK_ERROR(options.mM >= options.mTileM, "M must be equal or larger than TileM."); - TLLM_CHECK_ERROR(options.mM % options.mTileM == 0, "M must be divisible by TileM."); - TLLM_CHECK_ERROR(options.mTransposeMmaOutput, - "When batchN the MMA output has to be in column-major."); - } - - if (options.mUseDeepSeekFp8) { - if (batchM) { - // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. - TLLM_CHECK_ERROR( - options.mN % 128 == 0 && options.mValidN % 128 == 0, - "GEMM-N and validN must be a multiple of 128 when using DeepSeek Fp8. Found ", options.mN, - " and validN=", options.mValidN); - } else { - // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. - TLLM_CHECK_ERROR( - options.mM % 128 == 0 && options.mValidM % 128 == 0, - "GEMM-M and validM must be a multiple of 128 when using DeepSeek Fp8. Found ", options.mM, - " and validM=", options.mValidM); - } - // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. - TLLM_CHECK_ERROR(options.mK % 128 == 0 && options.mValidK % 128 == 0, - "GEMM-K and validK must be a multiple of 128 when using DeepSeek Fp8. Found ", - options.mK, " and validK=", options.mValidK); - - TLLM_CHECK_ERROR(options.mDtypeC != tg::Dtype::E2m1 && options.mDtypeA == tg::Dtype::E4m3 && - options.mDtypeB == tg::Dtype::E4m3, - "E2m1 is not supported with DeepSeek FP8"); - } - - if (options.mRouteSfsImpl.has_value() && options.mRouteSfsImpl.value() != options.mRouteImpl) { - TLLM_CHECK_ERROR( - (options.mRouteSfsImpl.value() == RouteImpl::Ldgsts || - options.mRouteSfsImpl.value() == RouteImpl::LdgPlusSts) && - options.mRouteImpl == RouteImpl::Tma, - "RouteSfsImpl must be equal to RouteImpl, or Ldgsts/LdgPlusSts, when RouteImpl is Tma"); - } else if (!options.mRouteSfsImpl.has_value()) { - if (updateOptions) { - options.mRouteSfsImpl = options.mRouteImpl; - } else { - TLLM_LOG_ERROR("RouteSfsImpl must be specified"); - return false; - } - } - - TLLM_CHECK_ERROR(options.mRouteImpl != RouteImpl::LdgPlusSts, - "LdgPlusSts does not support routing the tokens"); - - if (options.mRouteSfsImpl.has_value() && options.mRouteSfsImpl.value() == RouteImpl::LdgPlusSts) { - TLLM_CHECK_ERROR(options.mTileK <= 512 && options.mTileK >= 128, - "LdgPlusSts only supports 128 <= tileK <= 512"); - } - - if (batchM) { - if (options.mDtypeA == tg::Dtype::MxE2m1 && options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { - TLLM_CHECK_ERROR(doesRouteImplUseNoRoute(options.mRouteImpl), - "RouteAct is not supported with dtypeA = MxE2m1 and MxFp8Fp6Fp4."); - } - } else { - if (options.mDtypeB == tg::Dtype::MxE2m1 && options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { - TLLM_CHECK_ERROR(doesRouteImplUseNoRoute(options.mRouteImpl), - "RouteAct is not supported with dtypeB = MxE2m1 and MxFp8Fp6Fp4."); - } - } - - TLLM_CHECK_ERROR(options.mUseTmaStore, "Only TMA store is supported."); - if (batchM) { - TLLM_CHECK_ERROR(options.mLayoutA == gemm::MatrixLayout::MajorK, - "Activations must be in k-major format"); - } else { - TLLM_CHECK_ERROR(options.mLayoutB == gemm::MatrixLayout::MajorK, - "Activations must be in k-major format"); - } - - if (tg::mmaKindIsBlockFmt(options.mMmaKind) && !options.mUseDeepSeekFp8) { - if (!doesRouteImplUseNoRoute(options.mRouteImpl)) { - if (batchM) { - TLLM_CHECK_ERROR(options.mSfLayoutA == tg::SfLayout::Linear, - "Tokens need use SF linear layout when being routed"); - } else { - // Note: if B is cast from a non-block format to a block format, there are no SFs to load. - TLLM_CHECK_ERROR( - options.mSfLayoutB == tg::SfLayout::Linear || !tg::dtypeIsBlockFmt(options.mDtypeB), - "Tokens need use SF linear layout when being routed"); - } - } - - if (doesRouteImplUseTma(options.mRouteSfsImpl.value())) { - TLLM_CHECK_ERROR(!batchM, "UTMALDG.GATHER4 only supported for batch N."); - - if (tg::mmaKindIsBlockFmt(options.mMmaKind)) { - int const numEltsPerSfRoute = batchM ? options.mSfBlockSizeA : options.mSfBlockSizeB; - TLLM_CHECK_ERROR(options.mTileK % (numEltsPerSfRoute * 16) == 0, - "tileK needs to be a multiple of 16 * numEltsPerSf (", numEltsPerSfRoute, - ") = ", numEltsPerSfRoute * 16); - } - } - - if (!batchM || doesRouteImplUseNoRoute(options.mRouteImpl)) { - TLLM_CHECK_ERROR(options.mSfLayoutA == tg::SfLayout::R128c4, - "options.mSfLayoutA has to be tg::SfLayout::R128c4 when not being routed"); - } - } - - if (!gemm::isBiasTypeNone(options.mBiasType)) { - TLLM_CHECK_ERROR((gemm::isBiasTypeN(options.mBiasType) && - options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM) || - (gemm::isBiasTypeM(options.mBiasType) && - options.mBatchMode == BatchedGemmOptions::BatchMode::BatchN), - "BatchedGemm supports only per channel bias."); - } - - // We do not handle the case where K is not a multiple of TileK. - // TMA based load handles the case transparently. - if (doesRouteImplUseLdgsts(options.mRouteImpl) && - doesRouteImplUseLdgPlusSts(options.mRouteSfsImpl.value())) { - TLLM_CHECK_ERROR(options.mK % options.mTileK == 0, - "K must be a multiple of TileK when using Ldg based routing"); - } - - if (options.mRouteSfsImpl.has_value() && - (doesRouteImplUseLdgsts(options.mRouteSfsImpl.value()) || - doesRouteImplUseLdgPlusSts(options.mRouteSfsImpl.value()))) { - TLLM_CHECK_ERROR(options.mK % options.mTileK == 0, - "K must be a multiple of tileK when using Ldg based SF routing"); - } - - if (options.mClusterDimX > 1 && batchM && options.mRouteSfsImpl.has_value()) { - TLLM_CHECK_ERROR(options.mRouteSfsImpl.value() != RouteImpl::Tma, - "2CTA BatchedGemm does not support routing Sf along M dimension with TMA."); - } - - // Check if all elements in mBatchedM or mBatchedN are the same (uniform tokens per batch) and - // set mIsUniformNumTokensPerBatch and mBatchStride. - if (options.mIsUniformNumTokensPerBatch) { - int32_t firstValue = 0; - bool isUniformNumTokensPerBatch = false; - if (batchM && !options.mBatchedM.empty()) { - firstValue = options.mBatchedM[0]; - isUniformNumTokensPerBatch = std::all_of(options.mBatchedM.begin(), options.mBatchedM.end(), - [firstValue](int32_t v) { return v == firstValue; }); - } else if (!batchM && !options.mBatchedN.empty()) { - firstValue = options.mBatchedN[0]; - isUniformNumTokensPerBatch = std::all_of(options.mBatchedN.begin(), options.mBatchedN.end(), - [firstValue](int32_t v) { return v == firstValue; }); - } else { - TLLM_CHECK_ERROR( - false, "mBatchedM or mBatchedN must be specified when using uniform tokens per batch."); - } - auto tileTokensDim = batchM ? options.mTileM : options.mTileN; - TLLM_CHECK_ERROR(isUniformNumTokensPerBatch, - "All elements in mBatchedM or mBatchedN must be the same when using uniform " - "tokens per batch."); - TLLM_CHECK_ERROR(options.mBatchStrideInTokens >= 0, - "Batch stride in tokens must be greater or equal to 0 when using uniform " - "tokens per batch."); - TLLM_CHECK_ERROR_FMT( - options.mBatchStrideInTokens == 0 || - options.mBatchStrideInTokens == gemm::divUpMul(firstValue, tileTokensDim), - "Batch stride in tokens must be a 0 or a multiple of %s {%d} when using " - "uniform tokens per batch.", - batchM ? "TileM" : "TileN", tileTokensDim); - TLLM_CHECK_ERROR( - !options.mUseDeepSeekFp8, - "Uniform number of tokens per batch is not supported when using DeepSeek Fp8."); - TLLM_CHECK_ERROR( - !options.mUsePerTokenSfA && !options.mUsePerTokenSfB, - "Uniform number of tokens per batch is not supported when using per-token SF."); - TLLM_CHECK_ERROR(options.mBiasType == gemm::BiasType::None, - "Uniform number of tokens per batch is not supported when using bias."); - TLLM_CHECK_ERROR(options.mRouteImpl == RouteImpl::NoRoute, - "Uniform number of tokens per batch is not supported when using routing."); - TLLM_CHECK_ERROR( - !options.mFusedAct, - "Uniform number of tokens per batch is not supported when using fused gated activation."); - TLLM_CHECK_ERROR(!tg::dtypeIsBlockFmt(options.mDtypeA) && - !tg::dtypeIsBlockFmt(options.mDtypeB) && - !tg::dtypeIsBlockFmt(options.mDtypeC), - "Uniform number of tokens per batch is not supported when using block " - "format for dtypeA, dtypeB, or dtypeC."); - } else if (options.mBatchStrideInTokens >= 0) { - TLLM_LOG_WARNING("Batch stride in tokens is set to ", options.mBatchStrideInTokens, - " but it is not used when not using uniform tokens per batch."); - } - return isValid; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// BatchedGemmConfig -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -struct BatchedGemmConfig { - uint8_t const* mData{nullptr}; - uint32_t mSize{0}; - uint32_t mSharedMemSize{0}; - char const* mFunctionName{nullptr}; - uint32_t mNumThreadsPerCTA{0}; - char const* mHash{nullptr}; - - std::string mGenCfgJsonStr{""}; - char const* mExecPath{nullptr}; - trtllm::gen::CudaRunner* mCudaRunner{nullptr}; - trtllm::gen::GenCfg* mGenCfg{nullptr}; - int32_t mInstanceIdx{0}; - - BatchedGemmOptions mOptions; - tg::CudaArch mSm{tg::CudaArch::Sm100a}; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline std::string dumpOptions(BatchedGemmOptions const& options, bool dumpRuntimeParams = true) { - std::stringstream ss; - ss << gemmGatedAct::dumpOptions(options, dumpRuntimeParams) << ", "; - if (dumpRuntimeParams) { - ss << "mBatchedM={}," << std::endl; - ss << "mBatchedN={}," << std::endl; - } - ss << "mBatchMode=batchedGemm::BatchedGemmOptions::BatchMode(" - << static_cast(options.mBatchMode) << ")," << std::endl; - if (dumpRuntimeParams) { - ss << "mBatchStrideInTokens=" << options.mBatchStrideInTokens << "," << std::endl; - } - ss << "mFusedAct=" << options.mFusedAct << "," << std::endl; - ss << "mGridWaitForPrimaryRouting=" << options.mGridWaitForPrimaryRouting << "," << std::endl; - ss << "mIsStaticBatch=" << options.mIsStaticBatch << "," << std::endl; - ss << "mIsUniformNumTokensPerBatch=" << options.mIsUniformNumTokensPerBatch << "," << std::endl; - if (dumpRuntimeParams) { - ss << "mNumBatches=" << options.mNumBatches << "," << std::endl; - } - ss << "mNumRegsPerThreadLoadA=" << options.mNumRegsPerThreadLoadA << "," << std::endl; - ss << "mNumRegsPerThreadLoadB=" << options.mNumRegsPerThreadLoadB << "," << std::endl; - ss << "mNumRegsPerThreadLoadSfA=" << options.mNumRegsPerThreadLoadSfA << "," << std::endl; - ss << "mNumRegsPerThreadLoadSfB=" << options.mNumRegsPerThreadLoadSfB << "," << std::endl; - if (dumpRuntimeParams) { - ss << "mNumTokens=" << options.mNumTokens << "," << std::endl; - } - ss << "mNumWarpsLoadA=" << options.mNumWarpsLoadA << "," << std::endl; - ss << "mNumWarpsLoadB=" << options.mNumWarpsLoadB << "," << std::endl; - ss << "mNumWarpsLoadSfA=" << options.mNumWarpsLoadSfA << "," << std::endl; - ss << "mNumWarpsLoadSfB=" << options.mNumWarpsLoadSfB << "," << std::endl; - ss << "mRouteImpl=batchedGemm::RouteImpl(" << static_cast(options.mRouteImpl) << ")," - << std::endl; - ss << "mRouteSfsImpl={batchedGemm::RouteImpl(" - << static_cast(options.mRouteSfsImpl.value()) << ")}," << std::endl; - ss << "mUseTmaOobOpt=" << options.mUseTmaOobOpt << std::endl; - return ss.str(); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace batchedGemm - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -#ifdef TLLM_GEN_EXPORT_INTERFACE - -#undef TLLM_CHECK_ERROR -#undef TLLM_CHECK_ERROR_FMT -#undef TLLM_CHECK_WARNING -#undef TLLM_LOG_WARNING -#undef TLLM_LOG_INFO -#undef TLLM_LOG_ERROR - -#endif // TLLM_GEN_EXPORT_INTERFACE - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h deleted file mode 100644 index f0b63e674e..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/Enums.h +++ /dev/null @@ -1,162 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include - -namespace batchedGemm { - -namespace gemm { - -enum class AllReduceAlgo : uint32_t { - // Does not apply all-reduce. - None = 0, - // Reduction occurs at L2 cache; pulls N-1 partial outputs from peer devices. Result is - // non-deterministic. Potentially lower latency at cost of higher memory traffic. - OneShot, - // Reduction occurs at switch; pulls 1/Nth of the output from switch (reduce-scatter phase) and - // store to multicast mem (all-gather phase). Result is deterministic. Lower memory traffic at - // cost of potentially higher latency. - TwoShot, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -enum class MatrixLayout { - // K-major layout (default). [Mn, K] - MajorK = 0, - // M-major for A and N-major for B. [K, Mn] - MajorMn, - // Layout is blocked along the K dimension as seen in the diagram below. [K / blockK, Mn, blockK] - // where blockK is fixed at 128B - // - // ├────────────── K ──────────────┤ - // ┬ ┬ ├──── K block ───┤ - // │ │ │ 0 1 2 3 ║ 32 33 34 35 │ - // │ CTA0 │ 4 5 6 7 ║ 36 37 38 39 │ - // │ │ │ 8 9 10 11 ║ 40 41 42 43 │ - // │ ┴ │ 12 13 14 15 ║ 44 45 46 47 │ - // M ┬ ├────────────────║────────────────┤ - // │ │ │ 16 17 18 19 ║ 48 49 50 51 │ - // │ CTA1 │ 20 21 22 23 ║ 52 53 54 55 │ - // │ │ │ 24 25 26 27 ║ 56 57 58 59 │ - // ┴ ┴ │ 28 29 30 31 ║ 60 61 62 63 │ - BlockMajorK -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -enum class SplitK : uint32_t { - // No split-k is needed. I.e. mNumSlicesForSplitK == 1. - None = 0, - // CTAs computing one MN tile save partial results to global memory. - // Then wait on the barrier and the last CTA in the group loads partial results from gmem, - // sums them up and writes back to gmem. - Gmem, - // All CTAs in one CGA calculate partial sums. Then send the results to the smem of - // the last CTA in the CGA, which sums them up and writes to gmem. - Dsmem, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -enum class BiasType : uint32_t { - // No bias. - None = 0, - // One bias value per N of the output tensor. - M = 1, - // One bias value per row M of the output tensor. - N = 2, - // One bias value for each element of the output tensor. - Mn = 3, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Type of the element-wise activation to apply after the Gemm -enum class EltwiseActType { - None = 0, - // Gelu is defined as the following operation: - // act = x0 * phi(x0) - // where x0 is the output of the Gemm - // phi is the CDF of standard normal distribution approximated by - // phi(x) = 0.5 * (1 + tanh(0.7978845608028654 * (x + 0.044715 * x * x * x))) - Gelu, - // Relu2 (also known as squared Relu) is defined as the following operation: - // act = relu(x0) ^ 2 - // where x0 is the output of the Gemm. - Relu2, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -enum class TileScheduler { - // Static scheduler (Non-persistent). - Static = 0, - // Dynamic persistent scheduler. This is either based on an atomically incremented global work id - // prior to SM100 archs, or the HW supported work id scheduler based on UGETNEXTWORKID for SM100+. - Persistent, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -enum class CtaSwizzleType : uint32_t { - // Rasterize CTAs along the M dimension. - RasterizeAlongM = 0, - // Rasterize CTAs along the N dimension. - RasterizeAlongN, - // Swizzle CTAs in zig-zag pattern along M dimension, Zig-zag width is 2. - ZigZagAlongM2, - // Swizzle CTAs in zig-zag pattern along N dimension, Zig-zag width is 2. - ZigZagAlongN2, - // Swizzle CTAs in zig-zag pattern along M dimension, Zig-zag width is 4. - ZigZagAlongM4, - // Swizzle CTAs in zig-zag pattern along N dimension, Zig-zag width is 4. - ZigZagAlongN4, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Helper functions to check the SplitK type. - -#define SPLIT_K_FUNCTION(Mode) \ - inline bool doesSplitKUse##Mode(SplitK mode) { return (mode == SplitK::Mode); } - -SPLIT_K_FUNCTION(Gmem) -SPLIT_K_FUNCTION(Dsmem) - -#undef SPLIT_K_FUNCTION - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Helper functions to check the Bias type. - -#define BIAS_TYPE_FUNCTION(Mode) \ - inline bool isBiasType##Mode(BiasType type) { return (type == BiasType::Mode); } - -BIAS_TYPE_FUNCTION(None) -BIAS_TYPE_FUNCTION(N) -BIAS_TYPE_FUNCTION(M) -BIAS_TYPE_FUNCTION(Mn) - -#undef BIAS_TYPE_FUNCTION - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gemm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h deleted file mode 100644 index adca177a00..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmGatedActOptions.h +++ /dev/null @@ -1,236 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include "GemmOptions.h" - -#ifdef TLLM_GEN_EXPORT_INTERFACE -#include - -#define TLLM_CHECK_ERROR(cond, ...) \ - if (!(cond)) { \ - printArgs(__VA_ARGS__); \ - return false; \ - } - -#define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) - -#define TLLM_CHECK_ERROR_FMT(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) - -#define TLLM_CHECK_WARNING(cond, ...) \ - if (!(cond)) { \ - printArgs(__VA_ARGS__); \ - return false; \ - } - -#define TLLM_LOG_WARNING(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) - -#define TLLM_LOG_INFO(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) - -#endif - -namespace batchedGemm { - -namespace trtllm { -namespace gen { -class CudaRunner; -class GenCfg; -} // namespace gen -} // namespace trtllm - -namespace gemmGatedAct { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace tg = trtllm::gen; - -// Type of the gated activation -enum class ActType { - // clang-format off - // For ActType == SwiGlu, ideally we would like to have something like - // gatedAct = quantScaleC * (x0 * dequantScaleAb + beta) * ((x1 * scaleGate) * sigmoid(alpha * x1 * scaleGate)). - // But for now, we use the simplified version - // gatedAct = scaleC * (x0 + beta') * ((x1 * scaleGate) * sigmoid(alpha * x1 * scaleGate)), - // where x0 and x1 are the raw numbers from Gemm, while scaleC and scaleGate are input scales, - // beta' = beta / dequantScaleAb, scaleC = quantScaleC * dequantScaleAb. - // - // GatedSilu is a special case of SwiGlu where the alpha is 1.0 and the beta is 0.0. - // clang-format on - SwiGlu, - // For ActType == GeGlu, we use the simplified version - // gatedAct = scaleC' * (x0 + beta') * ((x1 * scaleGate) * phi(alpha * x1 * scaleGate)), - // where x0 and x1 are the raw numbers from Gemm, while scaleC and scaleGate are input scales, - // beta' = beta / scaleAb, scaleC' = scaleC * scaleAb. - GeGlu, - // Placeholder for no activation; not implemented in codegen - None, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Helper functions to check the ActType type. - -#define TLLM_ACT_TYPE_FUNCTION(actType) \ - inline bool is##actType(ActType type) { return (type == ActType::actType); } - -TLLM_ACT_TYPE_FUNCTION(SwiGlu) -TLLM_ACT_TYPE_FUNCTION(GeGlu) - -#undef TLLM_ACT_TYPE_FUNCTION - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline std::string getActTypeName(ActType type) { - switch (type) { - case ActType::SwiGlu: - return "SwiGlu"; - case ActType::GeGlu: - return "GeGlu"; - default: - return "Unknown type"; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -struct GemmGatedActOptions : public gemm::GemmOptions { - GemmGatedActOptions() = default; - GemmGatedActOptions(gemm::GemmOptions options, ActType actType, bool clampBeforeAct) - : gemm::GemmOptions(options), mActType(actType), mClampBeforeAct(clampBeforeAct) {} - - // Type of the gated activation. - ActType mActType{ActType::SwiGlu}; - // Clamp the dequantized values to the range [-limit, limit]. - bool mClampBeforeAct{false}; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Check if the options are valid or not. -inline bool checkAndUpdateGemmGatedActOptions(gemmGatedAct::GemmGatedActOptions& options, - tg::CudaArch cudaArch, bool updateOptions = true) { - auto isValid = gemm::checkAndUpdateGemmOptions(options, cudaArch, - /* tpGrpSize */ 1, updateOptions); - if (!isValid) { - return false; - } - - if (options.mActType == gemmGatedAct::ActType::None) { - TLLM_CHECK_ERROR(false, "ActType None is not supported"); - } - // tmpOut is already transposed at this stage - auto const hiddenSizeStr = options.mTransposeMmaOutput ? "M" : "N"; - auto const hiddenSize = options.mTransposeMmaOutput ? options.mM : options.mN; - auto const hiddenEpilogueTileSize = - options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN; - - TLLM_CHECK_ERROR(hiddenSize % 2 == 0, hiddenSizeStr, " must be a multiple of 2."); - - TLLM_CHECK_ERROR((options.mTransposeMmaOutput && !options.mUseShuffledMatrix) == false, - "Transpose mma output can only be used with shuffled matrix."); - - if (options.mUseTmaStore) { - TLLM_CHECK_ERROR( - hiddenEpilogueTileSize * tg::dtypeGetNumBits(options.mDtypeC) / /* bits */ 8 % 32 == 0, - "Unsupported output hidden tile size"); - } - - if (options.mDtypeC == tg::Dtype::E2m1 || options.mDtypeC == tg::Dtype::MxE4m3) { - int const outHiddenSize = (options.mTransposeMmaOutput ? options.mM : options.mN) / 2; - int const hiddenGranularity = 4 * options.mSfBlockSizeC; - TLLM_CHECK_ERROR(outHiddenSize % hiddenGranularity == 0, "Output hidden size (", outHiddenSize, - ") must be a multiple of ", hiddenGranularity, " for block-scaled outputs."); - } - - auto const validHiddenSize = options.mTransposeMmaOutput ? options.mValidM : options.mValidN; - if (options.mUseDeepSeekFp8) { - TLLM_CHECK_ERROR(hiddenSize % 256 == 0 && validHiddenSize % 256 == 0, "Hidden size (", - hiddenSize, ") and valid hidden size (", validHiddenSize, - ") must be a multiple of 256"); - } - - // - if (options.mUseShuffledMatrix) { - auto const shuffleBlockSize = gemm::getShuffleBlockSize(options.mEpilogueTileM); - TLLM_CHECK_ERROR( - hiddenSize % (2 * shuffleBlockSize) == 0 && validHiddenSize % (2 * shuffleBlockSize) == 0, - "M/validM must be a multiple of 2 * shuffle block size (", 2 * shuffleBlockSize, - ") when useShuffledMatrix"); - } - if (options.mNumSlicesForSplitK > 1) { - TLLM_CHECK_ERROR(doesSplitKUseDsmem(options.mSplitK), - "Split-k GMEM and GemmGatedAct are not supported yet."); - } - - if (gemm::isBiasTypeMn(options.mBiasType)) { - TLLM_CHECK_ERROR(options.mTransposeMmaOutput, - "Bias type Mn is not supported with not transpose mma output."); - } - - return true; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline std::string dumpOptions(GemmGatedActOptions const& options, bool dumpRuntimeParams = true) { - std::stringstream ss; - ss << gemm::dumpOptions(options, dumpRuntimeParams) << ", "; - ss << "mActType=" << "gemmGatedAct::ActType(" << static_cast(options.mActType) << ")," - << std::endl; - ss << "mClampBeforeAct=" << options.mClampBeforeAct << "" << std::endl; - return ss.str(); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// GemmGatedActConfig -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -struct GemmGatedActConfig { - uint8_t const* mData{nullptr}; - uint32_t mSize{0}; - uint32_t mSharedMemSize{0}; - char const* mFunctionName{nullptr}; - uint32_t mNumThreadsPerCTA{0}; - char const* mHash{nullptr}; - - std::string mGenCfgJsonStr{""}; - char const* mExecPath{nullptr}; - trtllm::gen::CudaRunner* mCudaRunner{nullptr}; - trtllm::gen::GenCfg* mGenCfg{nullptr}; - int32_t mInstanceIdx{0}; - - GemmGatedActOptions mOptions{}; - tg::CudaArch mSm{tg::CudaArch::Sm100a}; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gemmGatedAct - -#ifdef TLLM_GEN_EXPORT_INTERFACE - -#undef TLLM_CHECK_ERROR -#undef TLLM_CHECK_ERROR_FMT -#undef TLLM_CHECK_WARNING -#undef TLLM_LOG_WARNING -#undef TLLM_LOG_INFO -#undef TLLM_LOG_ERROR -#endif // TLLM_GEN_EXPORT_INTERFACE - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h deleted file mode 100644 index 1ee3fda469..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/GemmOptions.h +++ /dev/null @@ -1,1722 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include - -#include "Enums.h" -#include "KernelParams.h" -#include "KernelTraits.h" -#include "trtllm/gen/CudaArchDecl.h" -#include "trtllm/gen/DtypeDecl.h" -#include "trtllm/gen/MmaDecl.h" -#include "trtllm/gen/SfLayoutDecl.h" -#include "trtllm/gen/SparsityDecl.h" -#ifndef TLLM_GEN_EXPORT_INTERFACE -#include "trtllm/gen/CudaRunner.h" -#include "trtllm/gen/GenCtx.h" -#else -#ifdef TLLM_GEN_EXPORT_FLASHINFER -#include -namespace flashinfer::trtllm_cubin_loader { -std::string getCubin(const std::string& kernelName, const std::string& sha256); -} -#endif // TLLM_GEN_EXPORT_FLASHINFER -#include -namespace batchedGemm { - -template -void printArgs(T arg) { -#ifdef TLLM_GEN_DEBUG - std::cout << arg; -#endif -} - -template -void printArgs(T first, Args... args) { - printArgs(first); - if constexpr (sizeof...(args) > 0) { - printArgs(", "); - printArgs(args...); - } -} - -#define TLLM_CHECK_ERROR(cond, ...) \ - if (!(cond)) { \ - printArgs(__VA_ARGS__); \ - printArgs("\n"); \ - return false; \ - } - -#define TLLM_LOG_ERROR(...) TLLM_CHECK_ERROR(false, __VA_ARGS__) - -#define TLLM_CHECK_ERROR_FMT(cond, ...) TLLM_CHECK_ERROR(cond, __VA_ARGS__) - -#define TLLM_CHECK_WARNING(cond, ...) \ - if (!(cond)) { \ - printArgs(__VA_ARGS__); \ - printArgs("\n"); \ - return false; \ - } - -#define TLLM_LOG_WARNING(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) - -#define TLLM_LOG_INFO(...) TLLM_CHECK_WARNING(false, __VA_ARGS__) - -#endif // TLLM_GEN_EXPORT_INTERFACE - -#define GEMM_UPDATE_OR_ERROR(OPTION, VALUE) \ - if (updateOptions) { \ - OPTION = VALUE; \ - } else \ - return false - -namespace trtllm { -namespace gen { -class CudaRunner; -class GenCfg; -} // namespace gen -} // namespace trtllm - -namespace gemm { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace tg = trtllm::gen; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// NOTE: when adding new parameters, please update the dumpOptions function and -// gemm_export_config.json for cubin export. -struct GemmOptions { -#ifndef TLLM_GEN_EXPORT_INTERFACE - // allow safely down-casting - virtual ~GemmOptions() = default; -#endif - - GemmOptions() = default; - GemmOptions(AllReduceAlgo allReduceAlgo, BiasType biasType, int blockK, int clusterDimX, - int clusterDimY, int clusterDimZ, CtaSwizzleType ctaSwizzleType, tg::Dtype dtypeAcc, - tg::Dtype dtypeA, tg::Dtype dtypeB, tg::Dtype dtypeC, tg::Dtype dtypeMmaA, - tg::Dtype dtypeMmaB, EltwiseActType eltwiseActType, bool enablesEarlyExit, - bool enablesDelayedEarlyExit, bool enablesGlobalPtxKnobs, int epilogueLdtmDps, - int epilogueLdtmBits, int epilogueTileM, int epilogueTileN, bool fuseUtccpWithUtcmma, - bool gridTriggerSecondaryA, bool gridTriggerSecondaryB, - bool gridWaitForPrimaryEarlyExit, bool gridWaitForPrimaryA, bool gridWaitForPrimaryB, - bool hoistLoadTaskInit, bool hoistMmaTaskTryWaits, int k, KernelTraits kernelTraits, - MatrixLayout layoutA, MatrixLayout layoutB, int m, int mmaK, tg::MmaKind mmaKind, - int mmaM, int mmaN, bool mockAllReduce, int n, int numEpilogueWarps, - int numRegsCastAWarps, int numRegsCopySfLdsSttm, int numRegsCopySparsityInfo, - int numRegsPerThreadEpilogueWarp, int numRegsPerThreadNonEpilogueWarp, - int numSlicesForSplitK, int numSlicesForSliceK, int numStages, int numStagesMma, - int numStagesMmaWithinWorkTile, int numStagesMmaAcrossWorkTile, int numStagesWorkId, - bool outputDebugTensors, bool patchF2fp, int32_t sfBlockSizeA, int32_t sfBlockSizeB, - int32_t sfBlockSizeC, tg::SfLayout sfLayoutA, tg::SfLayout sfLayoutB, - tg::SfLayout sfLayoutC, int sfReshapeFactor, bool sliceK, tg::Sparsity sparsityA, - SplitK splitK, int tileK, int tileM, int tileN, TileScheduler tileScheduler, - bool transposeMmaOutput, bool useCustomMmaSchedule, bool useDeepSeekFp8, - bool useHoistTryWaitForCustomMmaSchedule, bool useMaxTmemOverlap, bool usePerTokenSfA, - bool usePerTokenSfB, bool useShuffledMatrix, bool useTmaStore, - bool useTwoTmaLoadWarps, bool useTwoMmaWarps, bool useUnrollLoop2xForMma, int validM, - int validN, int validK, int worldSize) - : mAllReduceAlgo{allReduceAlgo}, - mBiasType{biasType}, - mBlockK(blockK), - mClusterDimX{clusterDimX}, - mClusterDimY{clusterDimY}, - mClusterDimZ{clusterDimZ}, - mCtaSwizzleType{ctaSwizzleType}, - mDtypeAcc{dtypeAcc}, - mDtypeA{dtypeA}, - mDtypeB{dtypeB}, - mDtypeC{dtypeC}, - mDtypeMmaA{dtypeMmaA}, - mDtypeMmaB{dtypeMmaB}, - mEltwiseActType{eltwiseActType}, - mEnablesEarlyExit{enablesEarlyExit}, - mEnablesDelayedEarlyExit{enablesDelayedEarlyExit}, - mEnablesGlobalPtxKnobs{enablesGlobalPtxKnobs}, - mEpilogueLdtmDps{epilogueLdtmDps}, - mEpilogueLdtmBits{epilogueLdtmBits}, - mEpilogueTileM{epilogueTileM}, - mEpilogueTileN{epilogueTileN}, - mFuseUtccpWithUtcmma{fuseUtccpWithUtcmma}, - mGridTriggerSecondaryA{gridTriggerSecondaryA}, - mGridTriggerSecondaryB{gridTriggerSecondaryB}, - mGridWaitForPrimaryEarlyExit{gridWaitForPrimaryEarlyExit}, - mGridWaitForPrimaryA{gridWaitForPrimaryA}, - mGridWaitForPrimaryB{gridWaitForPrimaryB}, - mHoistLoadTaskInit{hoistLoadTaskInit}, - mHoistMmaTaskTryWaits{hoistMmaTaskTryWaits}, - mK{k}, - mKernelTraits{kernelTraits}, - mLayoutA{layoutA}, - mLayoutB{layoutB}, - mM{m}, - mMmaK{mmaK}, - mMmaKind{mmaKind}, - mMmaM{mmaM}, - mMmaN{mmaN}, - mMockAllReduce{mockAllReduce}, - mN{n}, - mNumEpilogueWarps{numEpilogueWarps}, - mNumRegsCastAWarps(numRegsCastAWarps), - mNumRegsCopySfLdsSttm(numRegsCopySfLdsSttm), - mNumRegsCopySparsityInfo(numRegsCopySparsityInfo), - mNumRegsPerThreadEpilogueWarp(numRegsPerThreadEpilogueWarp), - mNumRegsPerThreadNonEpilogueWarp(numRegsPerThreadNonEpilogueWarp), - mNumSlicesForSplitK{numSlicesForSplitK}, - mNumSlicesForSliceK{numSlicesForSliceK}, - mNumStages{numStages}, - mNumStagesMma{numStagesMma}, - mNumStagesMmaWithinWorkTile{numStagesMmaWithinWorkTile}, - mNumStagesMmaAcrossWorkTile{numStagesMmaAcrossWorkTile}, - mNumStagesWorkId{numStagesWorkId}, - mOutputDebugTensors{outputDebugTensors}, - mPatchF2fp{patchF2fp}, - mSfBlockSizeA{sfBlockSizeA}, - mSfBlockSizeB{sfBlockSizeB}, - mSfBlockSizeC{sfBlockSizeC}, - mSfLayoutA{sfLayoutA}, - mSfLayoutB{sfLayoutB}, - mSfLayoutC{sfLayoutC}, - mSfReshapeFactor{sfReshapeFactor}, - mSliceK{sliceK}, - mSparsityA{sparsityA}, - mSplitK{splitK}, - mTileK{tileK}, - mTileM{tileM}, - mTileN{tileN}, - mTileScheduler{tileScheduler}, - mTransposeMmaOutput{transposeMmaOutput}, - mUseCustomMmaSchedule{useCustomMmaSchedule}, - mUseDeepSeekFp8{useDeepSeekFp8}, - mUseHoistTryWaitForCustomMmaSchedule{useHoistTryWaitForCustomMmaSchedule}, - mUseMaxTmemOverlap{useMaxTmemOverlap}, - mUsePerTokenSfA{usePerTokenSfA}, - mUsePerTokenSfB{usePerTokenSfB}, - mUseShuffledMatrix{useShuffledMatrix}, - mUseTmaStore{useTmaStore}, - mUseTwoTmaLoadWarps{useTwoTmaLoadWarps}, - mUseTwoMmaWarps{useTwoMmaWarps}, - mUseUnrollLoop2xForMma{useUnrollLoop2xForMma}, - mValidM{validM}, - mValidN{validN}, - mValidK{validK}, - mWorldSize{worldSize} {} - // The all-reduce algorithm. - AllReduceAlgo mAllReduceAlgo{AllReduceAlgo::None}; - // The type of bias. - BiasType mBiasType{BiasType::None}; - // Block size in the K dimension - int mBlockK{-1}; - // Cluster size in X dim. - int mClusterDimX{1}; - // Cluster size in Y dim. - int mClusterDimY{1}; - // Cluster size in Z dim. - int mClusterDimZ{1}; - // The type of CTA swizzle. - CtaSwizzleType mCtaSwizzleType{CtaSwizzleType::RasterizeAlongM}; - // Data type of the accumulators. - tg::Dtype mDtypeAcc{tg::Dtype::Fp32}; - // Data type of the A matrix. - tg::Dtype mDtypeA{tg::Dtype::Fp16}; - // Data type of the B matrix. - tg::Dtype mDtypeB{tg::Dtype::Void}; - // Data type of the outputs. - tg::Dtype mDtypeC{tg::Dtype::Void}; - // Data type of the A matrix for the MMA, if different from the input type. - tg::Dtype mDtypeMmaA{tg::Dtype::Void}; - // Data type of the B matrix for the MMA, if different from the input type. - tg::Dtype mDtypeMmaB{tg::Dtype::Void}; - // The type of activation. - EltwiseActType mEltwiseActType{EltwiseActType::None}; - // Whether to enable early exit. - bool mEnablesEarlyExit{false}; - // Whether to enable delayed early exit to overlap - // numNonExitingCtas loading with the other instructions. - bool mEnablesDelayedEarlyExit{false}; - // Whether to enable the global PTX knobs for guiding the compiler optimizations. - bool mEnablesGlobalPtxKnobs{true}; - // The epilogue supports multiple LDTM shapes, although not every shape is applicable in every - // case. In particular: - // - On Hopper: must be 16dp256bit. - // - Transposed output: must be 16dp256bit. - // - Non-transposed output: - // - NvFp4 with fused activation: must be 32dp32bit. - // - Else it can be either 16dp256bit or 32dp32bit. - // The number of DP lanes in the epilogue LDTM. - int mEpilogueLdtmDps{16}; - // The number of bits in the epilogue LDTM. - int mEpilogueLdtmBits{256}; - // Tile size for the epilogue in M dimension. - int mEpilogueTileM{128}; - // Tile size for the epilogue in N dimension. - int mEpilogueTileN{32}; - // Whether fuse UTCCP with UTC*MMA. - bool mFuseUtccpWithUtcmma{false}; - // Whether load task A triggers the next grid. - bool mGridTriggerSecondaryA{false}; - // Whether load task B triggers the next grid. - bool mGridTriggerSecondaryB{false}; - // Whether the loads that check for an early exit should wait on a grid dependency. - bool mGridWaitForPrimaryEarlyExit{true}; - // Whether the load of A should wait on a grid dependency. - bool mGridWaitForPrimaryA{true}; - // Whether the load of B should wait on a grid dependency. - bool mGridWaitForPrimaryB{true}; - // Whether to hoist the initialization of the loading tasks. - bool mHoistLoadTaskInit{true}; - // Whether to hoist the mbarrier try_waits (e.g., mma.prodAcq, smemAb.consWait) in the MMA task. - bool mHoistMmaTaskTryWaits{false}; - // The K dimension of GEMM. - int mK{16 * 16}; - // Traits of the kernel. - KernelTraits mKernelTraits{}; - // Layout of A matrix - MatrixLayout mLayoutA{MatrixLayout::MajorK}; - // Layout of B matrix - MatrixLayout mLayoutB{MatrixLayout::MajorK}; - // The M dimension of GEMM. - int mM{128 * 2}; - // Size of the MMA instruction in the K dimension. - int mMmaK{16}; - // The kind of MMA instruction to use. - tg::MmaKind mMmaKind{tg::MmaKind::Auto}; - // Size of the MMA instruction in the M dimension. - int mMmaM{64}; - // Size of the MMA instruction in the N dimension. - int mMmaN{16}; - // Whether to mock all-reduce code for single-GPU debugging. - bool mMockAllReduce{false}; - // The N dimension of GEMM. - int mN{64 * 4}; - // Number of Epilogue Warps - int mNumEpilogueWarps{4}; - // Number of registers for the cast A warps. - int mNumRegsCastAWarps{0}; - // Number of registers for the LDS+STTM warps. - int mNumRegsCopySfLdsSttm{0}; - // Number of registers per thread to copy sparsity info with LDS+STTM. - int mNumRegsCopySparsityInfo{0}; - // Number of registers per thread for epilogue warps - int mNumRegsPerThreadEpilogueWarp{0}; - // Number of registers per thread for non-epilogue warps - int mNumRegsPerThreadNonEpilogueWarp{0}; - // Number of partitions along the K dimension. When mNumSlicesForSplitK > 1, - // the problem is distributed across several SMs, where each CTA works on its local K slice. - // Partial results are accumulated afterwards using either GMEM or DSMEM (in CGA) - // to exchange the data between CTAs. - int mNumSlicesForSplitK{1}; - // Number of slices for slice-K along K dimension. - int mNumSlicesForSliceK{1}; - // The depth of the mainloop pipeline. - int mNumStages{2}; - // The depth of the mma pipeline. Equals numStagesMmaWithinWorkTile * numStagesMmaAcrossWorkTile. - int mNumStagesMma{1}; - // The depth of the mma pipeline within work tile. Only GmemC classes with "WithAccInReg" suffix - // are allowed to be greater than 1. - int mNumStagesMmaWithinWorkTile{-1}; - // The depth of the mma pipeline across work tiles in the persistent loop. - int mNumStagesMmaAcrossWorkTile{-1}; - // The depth of the work id pipeline and the work throttle pipeline. - int mNumStagesWorkId{3}; - // Whether to output debug tensors. - bool mOutputDebugTensors{false}; - // Patch float conversions. - bool mPatchF2fp{false}; - // Block size of A, for block-scaled types. - int mSfBlockSizeA{-1}; - // Block size of B, for block-scaled types. - int mSfBlockSizeB{-1}; - // Block size of C, for block-scaled types. - int mSfBlockSizeC{-1}; - // Scale factors layout for A. - tg::SfLayout mSfLayoutA{tg::SfLayout::R128c4}; - // Scale factors layout for B. - tg::SfLayout mSfLayoutB{tg::SfLayout::R128c4}; - // Scale factors layout for C. - tg::SfLayout mSfLayoutC{tg::SfLayout::R128c4}; - // Number of "repeats", i.e. reshaping factor, to fold hidden dimension into SfBlock dimension. - // As result, the hidden dimension of the SF tensor must be a multiple of NumRepeats * - // numEltsPerSf * 4. This reduces the problem shape space that the kernel is able to run. - // But it reduces the number of L2 requests under the hood and potentially improves perf. - // Applies to layout 8x4 only. - int mSfReshapeFactor{1}; - // Slice-K implementation to use TileM dimension for TileK. - bool mSliceK{false}; - // Sparsity of A. - tg::Sparsity mSparsityA{tg::Sparsity::Dense}; - // The location of the exchange for split-K (it's None when split-K is disabled). - SplitK mSplitK{SplitK::None}; - // K tile dimension of GEMM. - int mTileK{16}; - // M tile dimension of GEMM. - int mTileM{128}; - // N tile dimension of GEMM. - int mTileN{32}; - // Tile scheduler type. - TileScheduler mTileScheduler{TileScheduler::Static}; - // Save output of MMA in M-major format. - bool mTransposeMmaOutput{false}; - // Use custom MMA schedule optimized for low-latency. - bool mUseCustomMmaSchedule{false}; - // Use DeepSeek Fp8. - bool mUseDeepSeekFp8{false}; - // The purpose of hoisting trywaits is to opportunistically peek at the availability of the next - // k-block. It benefits when the next k-block is already available and thus sustaining the - // momentum, but it adds latency to the first k-block for smaller k-loop. - bool mUseHoistTryWaitForCustomMmaSchedule{false}; - // Whether use the max Tmem overlap trick. - bool mUseMaxTmemOverlap{false}; - // Apply per-token scales from A - bool mUsePerTokenSfA{false}; - // Apply per-token scales from B - bool mUsePerTokenSfB{false}; - // Reorder rows/cols in the A matrix (when TransposeMmaOutput is true, otherwise B matrix) for the - // better memory accesses in the M-major epilogue. - bool mUseShuffledMatrix{false}; - // Use TMA to store the result. - bool mUseTmaStore{true}; - // Use two different warps for A and B matrix load. - bool mUseTwoTmaLoadWarps{false}; - // Use two different warps for MMA tasks. Applicable only to DeepSeek FP8. - bool mUseTwoMmaWarps{false}; - // Whether to unroll the loop by 2x. - bool mUseUnrollLoop2xForMma{true}; - // The valid range of M/N/K dimension of GEMM without padding values. - // Used to opportunistically remove memory traffic from the padding due to rigid SF shape - // constraint or TMA constraint. Such as: - // 1. outputDim % (4 * sfBlockSize) == 0; as 4x SFs are packed into 4 bytes - // 2. MxFp4 x Fp8 mmaType requires bespoke TMA load which requires hiddenDim % 128 == 0 - // 3. TMA requires 16B alignment for each row - int mValidM{-1}; - int mValidN{-1}; - int mValidK{-1}; - // World size for all-reduce. - int mWorldSize{1}; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -using SmVersion = tg::CudaArch; - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// GemmConfig -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -struct GemmConfig { - uint8_t const* mData{nullptr}; - uint32_t mSize{0}; - uint32_t mSharedMemSize{0}; - char const* mFunctionName{nullptr}; - uint32_t mNumThreadsPerCTA{0}; - char const* mHash{nullptr}; - std::string mGenCfgJsonStr{""}; - char const* mExecPath{nullptr}; - trtllm::gen::CudaRunner* mCudaRunner{nullptr}; - trtllm::gen::GenCfg* mGenCfg{nullptr}; - int32_t mInstanceIdx{0}; - - GemmOptions mOptions{}; - tg::CudaArch mSm{tg::CudaArch::Sm100a}; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Serialization helpers. -template -inline std::string toString(T e) { - return std::to_string(e); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline std::string toString(trtllm::gen::Dtype e) { - return trtllm::gen::dtypeToString(e); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template <> -inline std::string toString(trtllm::gen::MmaKind e) { - return trtllm::gen::mmaKindToString(e); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline std::string dumpOptions(GemmOptions const& options, bool dumpRuntimeParams = true) { - std::stringstream ss; - ss << "mAllReduceAlgo=" << "gemm::AllReduceAlgo(" << static_cast(options.mAllReduceAlgo) - << ")" << "," << std::endl; - ss << "mBiasType=" << "gemm::BiasType(" << static_cast(options.mBiasType) << ")" << "," - << std::endl; - ss << "mBlockK=" << options.mBlockK << "," << std::endl; - ss << "mClusterDimX=" << options.mClusterDimX << "," << std::endl; - ss << "mClusterDimY=" << options.mClusterDimY << "," << std::endl; - ss << "mClusterDimZ=" << options.mClusterDimZ << "," << std::endl; - ss << "mCtaSwizzleType=" << "gemm::CtaSwizzleType(" - << static_cast(options.mCtaSwizzleType) << ")" << "," << std::endl; - ss << "mDtypeAcc=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeAcc) << ")" - << "," << std::endl; - ss << "mDtypeA=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeA) << ")" << "," - << std::endl; - ss << "mDtypeB=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeB) << ")" << "," - << std::endl; - ss << "mDtypeC=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeC) << ")" << "," - << std::endl; - ss << "mDtypeMmaA=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeMmaA) << ")" - << "," << std::endl; - ss << "mDtypeMmaB=" << "trtllm::gen::Dtype(" << static_cast(options.mDtypeMmaB) << ")" - << "," << std::endl; - ss << "mEltwiseActType=" << "gemm::EltwiseActType(" - << static_cast(options.mEltwiseActType) << ")" << "," << std::endl; - ss << "mEnablesEarlyExit=" << options.mEnablesEarlyExit << "," << std::endl; - ss << "mEnablesDelayedEarlyExit=" << options.mEnablesDelayedEarlyExit << "," << std::endl; - ss << "mEnablesGlobalPtxKnobs=" << options.mEnablesGlobalPtxKnobs << "," << std::endl; - ss << "mEpilogueLdtmDps=" << options.mEpilogueLdtmDps << "," << std::endl; - ss << "mEpilogueLdtmBits=" << options.mEpilogueLdtmBits << "," << std::endl; - ss << "mEpilogueTileM=" << options.mEpilogueTileM << "," << std::endl; - ss << "mEpilogueTileN=" << options.mEpilogueTileN << "," << std::endl; - ss << "mFuseUtccpWithUtcmma=" << options.mFuseUtccpWithUtcmma << "," << std::endl; - ss << "mGridTriggerSecondaryA=" << options.mGridTriggerSecondaryA << "," << std::endl; - ss << "mGridTriggerSecondaryB=" << options.mGridTriggerSecondaryB << "," << std::endl; - ss << "mGridWaitForPrimaryEarlyExit=" << options.mGridWaitForPrimaryEarlyExit << "," << std::endl; - ss << "mGridWaitForPrimaryA=" << options.mGridWaitForPrimaryA << "," << std::endl; - ss << "mGridWaitForPrimaryB=" << options.mGridWaitForPrimaryB << "," << std::endl; - ss << "mHoistLoadTaskInit=" << options.mHoistLoadTaskInit << "," << std::endl; - ss << "mHoistMmaTaskTryWaits=" << options.mHoistMmaTaskTryWaits << "," << std::endl; - if (dumpRuntimeParams) { - ss << "mK=" << options.mK << "," << std::endl; - } - ss << "mKernelTraits={}" << "," << std::endl; - ss << "mLayoutA=gemm::MatrixLayout(" << static_cast(options.mLayoutA) << ")" << "," - << std::endl; - ss << "mLayoutB=gemm::MatrixLayout(" << static_cast(options.mLayoutB) << ")" << "," - << std::endl; - if (dumpRuntimeParams) { - ss << "mM=" << options.mM << "," << std::endl; - } - ss << "mMmaK=" << options.mMmaK << "," << std::endl; - ss << "mMmaKind=" << "trtllm::gen::MmaKind(" << static_cast(options.mMmaKind) << ")" - << "," << std::endl; - ss << "mMmaM=" << options.mMmaM << "," << std::endl; - ss << "mMmaN=" << options.mMmaN << "," << std::endl; - ss << "mMockAllReduce=" << options.mMockAllReduce << "," << std::endl; - if (dumpRuntimeParams) { - ss << "mN=" << options.mN << "," << std::endl; - } - ss << "mNumEpilogueWarps=" << options.mNumEpilogueWarps << "," << std::endl; - ss << "mNumRegsCastAWarps=" << options.mNumRegsCastAWarps << "," << std::endl; - ss << "mNumRegsCopySfLdsSttm=" << options.mNumRegsCopySfLdsSttm << "," << std::endl; - ss << "mNumRegsCopySparsityInfo=" << options.mNumRegsCopySparsityInfo << "," << std::endl; - ss << "mNumRegsPerThreadEpilogueWarp=" << options.mNumRegsPerThreadEpilogueWarp << "," - << std::endl; - ss << "mNumRegsPerThreadNonEpilogueWarp=" << options.mNumRegsPerThreadNonEpilogueWarp << "," - << std::endl; - ss << "mNumSlicesForSplitK=" << options.mNumSlicesForSplitK << "," << std::endl; - ss << "mNumSlicesForSliceK=" << options.mNumSlicesForSliceK << "," << std::endl; - ss << "mNumStages=" << options.mNumStages << "," << std::endl; - ss << "mNumStagesMma=" << options.mNumStagesMma << "," << std::endl; - ss << "mNumStagesMmaWithinWorkTile=" << options.mNumStagesMmaWithinWorkTile << "," << std::endl; - ss << "mNumStagesMmaAcrossWorkTile=" << options.mNumStagesMmaAcrossWorkTile << "," << std::endl; - ss << "mNumStagesWorkId=" << options.mNumStagesWorkId << "," << std::endl; - ss << "mOutputDebugTensors=" << options.mOutputDebugTensors << "," << std::endl; - ss << "mPatchF2fp=" << options.mPatchF2fp << "," << std::endl; - ss << "mSfBlockSizeA=" << options.mSfBlockSizeA << "," << std::endl; - ss << "mSfBlockSizeB=" << options.mSfBlockSizeB << "," << std::endl; - ss << "mSfBlockSizeC=" << options.mSfBlockSizeC << "," << std::endl; - ss << "mSfLayoutA=" << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutA) << ")" - << "," << std::endl; - ss << "mSfLayoutB=" << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutB) << ")" - << "," << std::endl; - ss << "mSfLayoutC=" << "trtllm::gen::SfLayout(" << static_cast(options.mSfLayoutC) << ")" - << "," << std::endl; - ss << "mSfReshapeFactor=" << options.mSfReshapeFactor << "," << std::endl; - ss << "mSliceK=" << options.mSliceK << "," << std::endl; - ss << "mSparsityA=" << "trtllm::gen::Sparsity(" << static_cast(options.mSparsityA) << ")" - << "," << std::endl; - ss << "mSplitK=" << "gemm::SplitK(" << static_cast(options.mSplitK) << ")" << "," - << std::endl; - ss << "mTileK=" << options.mTileK << "," << std::endl; - ss << "mTileM=" << options.mTileM << "," << std::endl; - ss << "mTileN=" << options.mTileN << "," << std::endl; - ss << "mTileScheduler=" << "gemm::TileScheduler(" << static_cast(options.mTileScheduler) - << ")" << "," << std::endl; - ss << "mTransposeMmaOutput=" << options.mTransposeMmaOutput << "," << std::endl; - ss << "mUseCustomMmaSchedule=" << options.mUseCustomMmaSchedule << "," << std::endl; - ss << "mUseDeepSeekFp8=" << options.mUseDeepSeekFp8 << "," << std::endl; - ss << "mUseHoistTryWaitForCustomMmaSchedule=" << options.mUseHoistTryWaitForCustomMmaSchedule - << "," << std::endl; - ss << "mUseMaxTmemOverlap=" << options.mUseMaxTmemOverlap << "," << std::endl; - ss << "mUsePerTokenSfA=" << options.mUsePerTokenSfA << "," << std::endl; - ss << "mUsePerTokenSfB=" << options.mUsePerTokenSfB << "," << std::endl; - ss << "mUseShuffledMatrix=" << options.mUseShuffledMatrix << "," << std::endl; - ss << "mUseTmaStore=" << options.mUseTmaStore << "," << std::endl; - ss << "mUseTwoTmaLoadWarps=" << options.mUseTwoTmaLoadWarps << "," << std::endl; - ss << "mUseTwoMmaWarps=" << options.mUseTwoMmaWarps << "," << std::endl; - ss << "mUseUnrollLoop2xForMma=" << options.mUseUnrollLoop2xForMma << "," << std::endl; - if (dumpRuntimeParams) { - ss << "mValidM=" << options.mValidM << "," << std::endl; - ss << "mValidN=" << options.mValidN << "," << std::endl; - ss << "mValidK=" << options.mValidK << "," << std::endl; - ss << "mWorldSize=" << options.mWorldSize << std::endl; - } - return ss.str(); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline T divUp(T a, T b) { - return (a + b - 1) / b; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline T divUpMul(T a, T b) { - return gemm::divUp(a, b) * b; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// clang-format off -inline std::vector srcToDstBlk16RowMap = - { - 0, 8, - 1, 9, - 2, 10, - 3, 11, - 4, 12, - 5, 13, - 6, 14, - 7, 15 - }; -inline std::vector srcToDstBlk32RowMap = - { - 0, 8, 16, 24, - 1, 9, 17, 25, - 2, 10, 18, 26, - 3, 11, 19, 27, - 4, 12, 20, 28, - 5, 13, 21, 29, - 6, 14, 22, 30, - 7, 15, 23, 31 - }; -// clang-format on - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getShuffleBlockSize(int epilogueTileM) { - int shuffleBlockSize = 16; - if (epilogueTileM % 128 == 0) { - shuffleBlockSize = 32; - } - return shuffleBlockSize; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline std::vector const& getShuffleIndices(int epilogueTileM) { - auto const shuffleBlockSize = getShuffleBlockSize(epilogueTileM); - return shuffleBlockSize == 16 ? srcToDstBlk16RowMap : srcToDstBlk32RowMap; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Check if the options are valid or not. -inline bool checkAndUpdateGemmOptions(GemmOptions& options, tg::CudaArch cudaArch, int tpGrpSize, - bool updateOptions = true) { - options.mWorldSize = tpGrpSize; - - bool isBlackwell = tg::isArchBlackwell(cudaArch); - - // If dtypeB is unspecified (Dtype::Void), assign to dtypeA. - if (options.mDtypeB == tg::Dtype::Void) { - if (updateOptions) { - options.mDtypeB = options.mDtypeA; - } else { - return false; - } - } - // If dtypeC is unspecified (Dtype::Void), assign to dtypeA. - if (options.mDtypeC == tg::Dtype::Void) { - TLLM_LOG_INFO("Setting dtypeC to ", tg::dtypeToString(options.mDtypeA)); - if (updateOptions) { - options.mDtypeC = options.mDtypeA; - } else { - return false; - } - } - - // If not specified, used the input dtypes as MMA dtypes (no cast required). - if (options.mDtypeMmaA == tg::Dtype::Void) { - if (updateOptions) { - options.mDtypeMmaA = options.mDtypeA; - } else { - return false; - } - } - if (options.mDtypeMmaB == tg::Dtype::Void) { - if (updateOptions) { - options.mDtypeMmaB = options.mDtypeB; - } else { - return false; - } - } - - // If validM/N/K is not specified, then assume the full range of the dimension is valid. - if (options.mValidM < 0 || options.mValidN < 0 || options.mValidK < 0) { - if (updateOptions) { - options.mValidM = options.mValidM < 0 ? options.mM : options.mValidM; - options.mValidN = options.mValidN < 0 ? options.mN : options.mValidN; - options.mValidK = options.mValidK < 0 ? options.mK : options.mValidK; - } else { - return false; - } - } - - // It must not exceed the padded dimensions. - if (options.mValidM > options.mM || options.mValidN > options.mN || - options.mValidK > options.mK) { - TLLM_LOG_WARNING( - "ValidM, ValidN, and ValidK must be less than or equal to M, N, and K respectively."); - if (updateOptions) { - options.mValidM = std::min(options.mValidM, options.mM); - options.mValidN = std::min(options.mValidN, options.mN); - options.mValidK = std::min(options.mValidK, options.mK); - } else { - return false; - } - } - - // BlockMajorK layout does not support validM, validN, validK parameters - if (options.mLayoutA == gemm::MatrixLayout::BlockMajorK || - options.mLayoutB == gemm::MatrixLayout::BlockMajorK) { - bool hasValidParams = (options.mValidM != -1 && options.mValidM != options.mM) || - (options.mValidN != -1 && options.mValidN != options.mN) || - (options.mValidK != -1 && options.mValidK != options.mK); - TLLM_CHECK_ERROR(!hasValidParams, - "BlockMajorK layout does not support validM/validN/validK parameters due to " - "swizzled layout. " - "Found validM=", - options.mValidM, " validN=", options.mValidN, " validK=", options.mValidK); - } - -#ifdef TLLM_PUBLIC_RELEASE - if (options.mDtypeA == tg::Dtype::E2m1 && options.mDtypeMmaA == tg::Dtype::E4m3) { - TLLM_CHECK_ERROR(false, "E2m1 x E4m3 is not supported for JIT compile. Use cubins instead."); - } -#endif // TLLM_PUBLIC_RELEASE - - // Check that the A cast is supported. - // Currently, we only support {MxFp4, NvFp4, MxInt4} -> Bf16. - TLLM_CHECK_ERROR( - (options.mDtypeA == options.mDtypeMmaA) || - ((options.mDtypeA == tg::Dtype::MxE2m1 || options.mDtypeA == tg::Dtype::E2m1 || - options.mDtypeA == tg::Dtype::MxInt4) && - options.mDtypeMmaA == tg::Dtype::Bfloat16) || - (options.mDtypeA == tg::Dtype::E2m1 && options.mDtypeMmaA == tg::Dtype::E4m3), - "Unsupported cast for A: ", tg::dtypeToString(options.mDtypeA), " -> ", - tg::dtypeToString(options.mDtypeMmaA)); - - // Check that the B cast is supported. - // Currently, we only support Fp8 -> MxFp8. - // TODO: add same support for A (no transpose) - TLLM_CHECK_ERROR( - (options.mDtypeB == options.mDtypeMmaB) || - (options.mDtypeB == tg::Dtype::E4m3 && options.mDtypeMmaB == tg::Dtype::MxE4m3), - "Unsupported cast for B: ", tg::dtypeToString(options.mDtypeB), " -> ", - tg::dtypeToString(options.mDtypeMmaB)); - - if (options.mDtypeA != options.mDtypeMmaA) { - TLLM_CHECK_ERROR(options.mTileM == 128, - "TileM must be 128 when casting the input matrix A before the MMA."); - } - - if (options.mPatchF2fp) { - TLLM_CHECK_ERROR( - options.mDtypeA == tg::Dtype::MxE2m1 && options.mDtypeMmaA == tg::Dtype::Bfloat16, - "PatchF2fp is only supported for MxFp4 to Bf16 casts."); - } -#ifdef TLLM_PUBLIC_RELEASE - options.mPatchF2fp = false; -#endif // TLLM_PUBLIC_RELEASE - - // FIXME: We do not support different dtypes for A and B when not on Blackwell. - if (!isBlackwell) { - TLLM_CHECK_ERROR(options.mDtypeMmaA == options.mDtypeMmaB, - "For non-Blackwell, A and B must have the same dtype."); - } - - // Check that the different dtypes for A and B are supported by the tensor core - // kind::f8f6f4 - if (options.mDtypeMmaA == tg::Dtype::E4m3 || options.mDtypeMmaA == tg::Dtype::E2m1) { - TLLM_CHECK_ERROR(options.mDtypeMmaB == tg::Dtype::E4m3 || options.mDtypeMmaB == tg::Dtype::E2m1, - "For dtypeMmaA = E4m3/E2m1 A, dtypeMmaB must also be E4m3/E2m1."); - } - - // kind::mxf8f6f4 - if (options.mDtypeMmaA == tg::Dtype::MxE4m3 || options.mDtypeMmaA == tg::Dtype::MxE2m1) { - TLLM_CHECK_ERROR( - options.mDtypeMmaB == tg::Dtype::MxE4m3 || options.mDtypeMmaB == tg::Dtype::MxE2m1, - "For dtypeMmaA = MxE4m3 or MxE2m1, dtypeMmaB must also be MxE4m3 or MxE2m1."); - } - if (options.mDtypeMmaB == tg::Dtype::MxE4m3 || options.mDtypeMmaB == tg::Dtype::MxE2m1) { - TLLM_CHECK_ERROR( - options.mDtypeMmaA == tg::Dtype::MxE4m3 || options.mDtypeMmaA == tg::Dtype::MxE2m1, - "For dtypeMmaB = MxE4m3 or MxE2m1, dtypeMmaA must also be MxE4m3 or MxE2m1."); - } - - // kind::f16 - if (options.mDtypeMmaA == tg::Dtype::Fp16 || options.mDtypeMmaA == tg::Dtype::Bfloat16) { - TLLM_CHECK_ERROR(options.mDtypeMmaB == options.mDtypeMmaA, - "For dtypeMmaA = Fp16/Bfloat16, dtypeMmaB must be the same as dtypeMmaA."); - } - - // When one of the inputs needs to be cast, we must use two load warps. - if ((options.mDtypeMmaA != options.mDtypeA || options.mDtypeMmaB != options.mDtypeB) && - !options.mUseTwoTmaLoadWarps) { - TLLM_LOG_WARNING("Two TMA load warps must be enabled if any of the inputs needs to be cast."); - } - - // When different dtypes are used for A and B, we must use different tiles to do the loading. - // It is not strictly required, but current implementation of SmemAb requires that. - if (options.mDtypeA != options.mDtypeB) { - TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps, - "Two TMA load warps must be enabled for different input types of A and B."); - } - - // Get the mma kind for the input types. - if (options.mMmaKind == tg::MmaKind::Auto) { - if (updateOptions) { - options.mMmaKind = dtypeGetMmaKind(options.mDtypeMmaA, options.mDtypeMmaB); - } else { - return false; - } - } - - // Check that the sparsity mode of A is supported, and compatible with the MMA kind. - // Note: trtllm-gen currently does not support sparsity with tf32, fp16, bf16. - switch (options.mSparsityA) { - case tg::Sparsity::Dense: - // Always supported. - break; - case tg::Sparsity::Any_1_2: - TLLM_LOG_ERROR("1:2 sparsity is not supported."); - break; - case tg::Sparsity::Any_2_4: { - bool isSupported_2_4 = (options.mMmaKind == tg::MmaKind::Fp8Fp6Fp4 || - options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4); - TLLM_CHECK_ERROR(isSupported_2_4, "2:4 sparsity is not supported for MMA kind ", - tg::mmaKindToString(options.mMmaKind), " on target ", - tg::cudaArchToString(cudaArch)); - break; - } - case tg::Sparsity::Pairwise_4_8: - TLLM_CHECK_ERROR(options.mMmaKind == tg::MmaKind::MxFp4NvFp4, - "Pairwise 4:8 sparsity is only supported for MMA kind MxFp4NvFp4."); - break; - default: - TLLM_CHECK_ERROR(false, "Unsupported sparsityA: ", tg::sparsityToString(options.mSparsityA)); - break; - } - - // Is A sparse? - bool const isSparseA = tg::isSparse(options.mSparsityA); - - // Requirements specific to sparsity, and compatibility with other features. - if (isSparseA) { - TLLM_CHECK_ERROR(isBlackwell, "Sparsity is only supported on Blackwell"); - // The following requirement is for TMA load: the box width must be a multiple of 16B. - TLLM_CHECK_ERROR( - tg::getNumBytesSparsityInfo(options.mSparsityA, options.mTileK) % 16 == 0, - "The sparsity information for one tile row must be a multiple of 16B. Use larger tileK."); - TLLM_CHECK_ERROR(options.mDtypeA == options.mDtypeMmaA, - "Sparsity is not supported with on-the-fly upcasting."); - TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, "Sparsity is not supported with DeepSeek Fp8."); - TLLM_CHECK_ERROR(!options.mSliceK, "Sparsity is not supported with slice-k."); - } - - if (options.mMmaKind == tg::MmaKind::Fp8Fp6Fp4) { - int mmaK = isSparseA ? 64 : 32; - - if (options.mMmaK != mmaK) { - TLLM_LOG_WARNING( - "Unsupported MmaK (", options.mMmaK, ") for MmaKind=", gemm::toString(options.mMmaKind), - " and sparsity=", tg::sparsityToString(options.mSparsityA), ". Setting MmaK to ", mmaK); - if (updateOptions) { - options.mMmaK = mmaK; - options.mTileK = std::max(options.mMmaK, options.mTileK); - } else { - return false; - } - } - } - - // Check LDTM shape. - if (isBlackwell) { - TLLM_CHECK_ERROR((options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256) || - (options.mEpilogueLdtmDps == 32 && options.mEpilogueLdtmBits == 32), - "Unsupported LDTM shape: ", options.mEpilogueLdtmDps, "dp", - options.mEpilogueLdtmBits, "bit."); - if (options.mEpilogueTileM == 64) { - TLLM_CHECK_ERROR(options.mEpilogueLdtmDps == 16, - "Unsupported LDTM shape for epilogueTileM=64: ", options.mEpilogueLdtmDps, - "dp", options.mEpilogueLdtmBits, "bit."); - } - if (options.mTransposeMmaOutput) { - // We can't use 32dp32bit LDTM for transposed outputs because we need each thread to own - // multiple consecutive output elements. - TLLM_CHECK_ERROR((options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256), - "Only 16dp256bit LDTM is supported for transposed outputs."); - } - } else { - TLLM_CHECK_ERROR( - options.mEpilogueLdtmDps == 16 && options.mEpilogueLdtmBits == 256, - "Hopper does not use TMEM. The register layout corresponds to 16dp256bit. Got ", - options.mEpilogueLdtmDps, "dp", options.mEpilogueLdtmBits, "bit."); - } - - // Constraints for NvFp4, MxFp8, and MxFp4. - if ((options.mMmaKind == tg::MmaKind::MxFp4NvFp4 || - options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4 || options.mDtypeC == tg::Dtype::MxE4m3) && - options.mMmaM != 128) { - if (options.mClusterDimX == 1) { - // MMA M must be 128 when the input uses block scaling, or when the output is an Mx format. - int newTileM = 128 * divUp(options.mTileM, 128); - TLLM_LOG_WARNING("Unsupported MmaM (", options.mMmaM, - ") for MmaKind=", gemm::toString(options.mMmaKind), - ". Setting MmaM to 128 and TileM to ", newTileM); - if (updateOptions) { - options.mMmaM = 128; - options.mTileM = newTileM; - } else { - return false; - } - } else { - TLLM_CHECK_ERROR(options.mMmaM == 256 && options.mTileM == 128, - "2CTA UTCxMMA only supports mmaM = 256 and tileM = 128."); - } - } - if (options.mClusterDimX > 1) { - TLLM_CHECK_ERROR(options.mLayoutB != MatrixLayout::BlockMajorK, - "layoutB == MatrixLayout::BlockMajorK is not supported for now"); - } - if (options.mMmaKind == tg::MmaKind::MxFp4NvFp4 || options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { - TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); - - int mmaK = isSparseA ? 64 : 32; - if (options.mMmaKind == tg::MmaKind::MxFp4NvFp4) { - mmaK = isSparseA ? 128 : 64; - if (options.mMmaK == 96 && !isSparseA) { - mmaK = 96; - TLLM_CHECK_ERROR(options.mTileK == 768, "When mmaK == 96, only tileK == 768 is supported"); - TLLM_CHECK_ERROR(options.mTileN <= 128, "When mmaK == 96, only tileN <= 128 is supported"); - } - } - if (options.mMmaK != mmaK) { - int newTileK = mmaK * divUp(options.mTileK, mmaK); - TLLM_LOG_WARNING("Unsupported MmaK (", options.mMmaK, - ") for MmaKind=", gemm::toString(options.mMmaKind), ". Setting MmaK to ", - mmaK, " and TileK to ", newTileK); - if (updateOptions) { - options.mMmaK = mmaK; - options.mTileK = newTileK; - } else { - return false; - } - } - - // The MMA N may only be smaller than 64 if it is equal to the tile N. - TLLM_CHECK_ERROR(options.mMmaN >= 64 || options.mMmaN == options.mTileN, "MmaN (", - options.mMmaN, ") must be >= 64 or equal to TileN (", options.mTileN, ")"); - } - - // Note: the logic for selecting/checking the correct block size based on dtypes and sparsity is - // centralized here, to avoid error-prone code duplication and make it a more explicit "contract" - // with the user who is providing inputs in this format. - // Additionally, in some cases, multiple values are possible: - // - When we use type casting before the MMA (e.g. e2m1 x e4m3). - // - For output C, based on whether the consumer will use sparsity. - - // SF block size for A. - if (options.mDtypeA == tg::Dtype::E2m1 && options.mDtypeB == tg::Dtype::E4m3) { - // Note that the type conversion needs to happen before TCs. - // For example, convert e2m1 to e4m3 inside TmemCastA. - if (!(options.mSfBlockSizeA == 16 || options.mSfBlockSizeA == 32)) { - TLLM_LOG_WARNING("sfBlockSizeA must be 16 or 32 for e2m1 x e4m3, got ", - options.mSfBlockSizeA); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeA, 16); - } - } else if (options.mDtypeA == tg::Dtype::E2m1) { - if (!((options.mSfBlockSizeA == 16 && !isSparseA) || - (options.mSfBlockSizeA == 32 && isSparseA))) { - TLLM_LOG_WARNING("sfBlockSizeA must be 16 (dense) or 32 (sparse) for dtypeA=e2m1, got ", - options.mSfBlockSizeA); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeA, isSparseA ? 32 : 16); - } - } else if (options.mDtypeA == tg::Dtype::MxE2m1 || options.mDtypeA == tg::Dtype::MxE4m3 || - options.mDtypeA == tg::Dtype::MxInt4) { - if (!((options.mSfBlockSizeA == 32 && !isSparseA) || - (options.mSfBlockSizeA == 64 && isSparseA))) { - TLLM_LOG_WARNING( - "sfBlockSizeA must be 32 (dense) or 64 (sparse) for dtypeA=mx{e2m1,e4m3,int4}, got ", - options.mSfBlockSizeA); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeA, isSparseA ? 64 : 32); - } - } else if (options.mSfBlockSizeA > 0) { - TLLM_LOG_WARNING("Got sfBlockSizeA=", options.mSfBlockSizeA, - " but dtypeA=", tg::dtypeToString(options.mDtypeA), - " does not use block scales"); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeA, -1); - } - // SF block size for B. - if (options.mDtypeB == tg::Dtype::E2m1) { - if (!((options.mSfBlockSizeB == 16 && !isSparseA) || - (options.mSfBlockSizeB == 32 && isSparseA))) { - TLLM_LOG_WARNING("sfBlockSizeB must be 16 (dense) or 32 (sparse) for dtypeB=e2m1, got ", - options.mSfBlockSizeB); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeB, isSparseA ? 32 : 16); - } - } else if (options.mDtypeB == tg::Dtype::MxE2m1 || options.mDtypeB == tg::Dtype::MxE4m3 || - (options.mDtypeB == tg::Dtype::E4m3 && options.mDtypeMmaB == tg::Dtype::MxE4m3)) { - if (!((options.mSfBlockSizeB == 32 && !isSparseA) || - (options.mSfBlockSizeB == 64 && isSparseA))) { - TLLM_LOG_WARNING( - "sfBlockSizeB must be 32 (dense) or 64 (sparse) for dtypeB=mx{e2m1,e4m3}, got ", - options.mSfBlockSizeB); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeB, isSparseA ? 64 : 32); - } - } else if (options.mSfBlockSizeB > 0) { - TLLM_LOG_WARNING("Got sfBlockSizeB=", options.mSfBlockSizeB, - " but dtypeB=", tg::dtypeToString(options.mDtypeB), - " does not use block scales"); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeB, -1); - } - // SF block size for C. - if (options.mDtypeC == tg::Dtype::E2m1) { - if (!(options.mSfBlockSizeC == 16 || options.mSfBlockSizeC == 32)) { - TLLM_LOG_WARNING("sfBlockSizeC must be 16 or 32 for dtypeC=e2m1, got ", - options.mSfBlockSizeC); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeC, 16); - } - } else if (options.mDtypeC == tg::Dtype::MxE2m1 || options.mDtypeC == tg::Dtype::MxE4m3) { - if (!(options.mSfBlockSizeC == 32 || options.mSfBlockSizeC == 64)) { - TLLM_LOG_WARNING("sfBlockSizeC must be 32 or 64 for dtypeC=mx{e2m1,e4m3}, got ", - options.mSfBlockSizeC); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeC, 32); - } - } else if (options.mSfBlockSizeC > 0) { - TLLM_LOG_WARNING("Got sfBlockSizeC=", options.mSfBlockSizeC, - " but dtypeC=", tg::dtypeToString(options.mDtypeC), - " does not use block scales"); - GEMM_UPDATE_OR_ERROR(options.mSfBlockSizeC, -1); - } - - if (tg::dtypeIsBlockFmt(options.mDtypeA)) { - int numEltsPerSfA = options.mSfBlockSizeA; - TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSfA) == 0, "TileK (", options.mTileK, - ") must be a multiple of ", (4 * numEltsPerSfA), " for typeA ", - gemm::toString(options.mDtypeA)); - auto const numEltsPerSfAInK = options.mK / numEltsPerSfA; - TLLM_CHECK_ERROR(numEltsPerSfAInK % 4 == 0, "K dimension of scaling factors for A (", - numEltsPerSfAInK, ") must be a multiple of 4"); - } - if (tg::dtypeIsBlockFmt(options.mDtypeB)) { - TLLM_CHECK_ERROR(options.mSfLayoutB == tg::SfLayout::R128c4 || - options.mSfLayoutB == tg::SfLayout::R8c4 || - options.mSfLayoutB == tg::SfLayout::Linear, - "Only the 128x4 and 8x4 SF layouts are supported for B, got ", - tg::sfLayoutToString(options.mSfLayoutB)); - - // TileN must be a multiple of the number of rows per SF tile. - int const numSfTileRowsB = options.mSfLayoutB == tg::SfLayout::R128c4 ? 128 : 8; - TLLM_CHECK_ERROR(options.mTileN % numSfTileRowsB == 0, "TileN (", options.mTileN, - ") must be a multiple of ", numSfTileRowsB, " for B SF layout ", - tg::sfLayoutToString(options.mSfLayoutB)); - - int numEltsPerSfB = options.mSfBlockSizeB; - TLLM_CHECK_ERROR(options.mTileK % (4 * numEltsPerSfB) == 0, "TileK (", options.mTileK, - ") must be a multiple of ", (4 * numEltsPerSfB), " for typeB ", - gemm::toString(options.mDtypeB)); - auto const numEltsPerSfBInK = options.mK / numEltsPerSfB; - TLLM_CHECK_ERROR(numEltsPerSfBInK % 4 == 0, "K dimension of scaling factors for B (", - numEltsPerSfBInK, ") must be a multiple of 4"); - } - - int32_t padMultiplierA = 1; - int32_t padMultiplierB = 1; - if (options.mMmaKind == tg::MmaKind::MxFp8Fp6Fp4) { - if (options.mDtypeA == tg::Dtype::MxE2m1) { - padMultiplierA = 2; - } - if (options.mDtypeB == tg::Dtype::MxE2m1) { - padMultiplierB = 2; - } - } - TLLM_CHECK_ERROR( - (padMultiplierA * tg::dtypeGetNumBits(options.mDtypeA) * options.mK / 8) % 16 == 0, - "K dimension of A must be aligned to 16 bytes."); - TLLM_CHECK_ERROR( - (padMultiplierB * tg::dtypeGetNumBits(options.mDtypeB) * options.mK / 8) % 16 == 0, - "K dimension of B must be aligned to 16 bytes."); - - if (tg::dtypeIsBlockFmt(options.mDtypeC)) { - TLLM_CHECK_ERROR(isBlackwell, "Block scaling is only supported on Blackwell"); - - TLLM_CHECK_ERROR( - options.mSfLayoutC == tg::SfLayout::R128c4 || options.mSfLayoutC == tg::SfLayout::R8c4, - "Only the 128x4 and 8x4 SF layouts are supported for C."); - if (!options.mTransposeMmaOutput) { - TLLM_CHECK_ERROR(options.mEpilogueTileN % options.mSfBlockSizeC == 0, - "EpilogueTileN must be a multiple of the number of elements per SF for C"); - } - int const numSfTileRowsC = options.mSfLayoutC == tg::SfLayout::R128c4 ? 128 : 8; - int const tileTokenDim = options.mTransposeMmaOutput ? options.mTileN : options.mTileM; - TLLM_CHECK_ERROR_FMT(tileTokenDim % numSfTileRowsC == 0, - "Tile%s (%d) must be a multiple of %d for C SF layout %s", - options.mTransposeMmaOutput ? "N" : "M", tileTokenDim, numSfTileRowsC, - tg::sfLayoutToString(options.mSfLayoutC).c_str()); - - int numEltsPerSfC = options.mSfBlockSizeC; - int const hiddenDim = options.mTransposeMmaOutput ? options.mM : options.mN; - int const hiddenGranularity = 4 * numEltsPerSfC; - TLLM_CHECK_ERROR(hiddenDim % hiddenGranularity == 0, "Hidden dim (", hiddenDim, - ") must be a multiple of ", hiddenGranularity, " for block-scaled outputs."); - int const validHiddenDim = options.mTransposeMmaOutput ? options.mValidM : options.mValidN; - TLLM_CHECK_ERROR(validHiddenDim % numEltsPerSfC == 0, "Valid hidden dim (", validHiddenDim, - ") must be a multiple of ", numEltsPerSfC, " for block-scaled outputs."); - TLLM_CHECK_ERROR(!options.mTransposeMmaOutput || options.mUseShuffledMatrix, - "Transposing block-scaled outputs requires shuffled matrix."); - } - - // Set epilogue tile sizes to the output tile sizes, when epilogue tile sizes are incorrect. - if (options.mTileM % options.mEpilogueTileM != 0) { - TLLM_LOG_WARNING("TileM (", options.mTileM, ") must be divisible by EpilogueTileM (", - options.mEpilogueTileM, "). Setting EpilogueTileM to TileM"); - if (updateOptions) { - options.mEpilogueTileM = options.mTileM; - } else { - return false; - } - } - - if (options.mTileN % options.mEpilogueTileN != 0) { - TLLM_LOG_WARNING("TileN (", options.mTileN, ") must be divisible by EpilogueTileN (", - options.mEpilogueTileN, "). Setting EpilogueTileN to TileN"); - if (updateOptions) { - options.mEpilogueTileN = options.mTileN; - } else { - return false; - } - } - - // On Hopper, epilogue tile sizes are the same as output tiles. - if (!isBlackwell && - (options.mEpilogueTileM != options.mTileM || options.mEpilogueTileN != options.mTileN)) { - TLLM_LOG_WARNING( - "Overwriting epilogueTileM and epilogueTileN to match tileM and tileN respectively"); - if (updateOptions) { - options.mEpilogueTileM = options.mTileM; - options.mEpilogueTileN = options.mTileN; - } else { - return false; - } - } - - // Unsupported epilogue tile size. - if (options.mMmaM == 128 && options.mEpilogueTileM != options.mTileM) { - TLLM_LOG_WARNING( - "When MmaM = 128, EpilogueTileM must be equal to TileM. Setting EpilogueTileM to TileM"); - if (updateOptions) { - options.mEpilogueTileM = options.mTileM; - } else { - return false; - } - } - - TLLM_CHECK_ERROR(options.mM > 0 && options.mN > 0 && options.mK > 0, - "M, N and K must be larger than 0"); - TLLM_CHECK_ERROR(options.mNumSlicesForSplitK > 0, "Split K must be larger than 0."); - - if (options.mUseShuffledMatrix) { - auto const shuffleBlockSize = getShuffleBlockSize(options.mEpilogueTileM); - TLLM_CHECK_ERROR(options.mM % shuffleBlockSize == 0 && options.mValidM % shuffleBlockSize == 0, - "M/validM must be a multiple of shuffle block size (", shuffleBlockSize, - ") when useShuffledMatrix"); - } - - if (!options.mSliceK) { - TLLM_CHECK_ERROR(options.mMmaM / options.mClusterDimX <= options.mEpilogueTileM, - "EpilogueTileM must be larger or equal than mmaM."); - } else { - // FIXME: this is not necessary limitation. Simply fixing num repeats in TmemSliceKA should be - // enough. - TLLM_CHECK_ERROR((options.mTileN & (options.mTileN - 1)) == 0, - "For Slice-K TileN is required to be a power of 2"); - } - - if (options.mClusterDimX == 2) { - TLLM_CHECK_ERROR(options.mMmaM == 256, "Only mmaM = 256 is supported for 2CTA UTCMMA."); - TLLM_CHECK_ERROR(options.mMmaN % 16 == 0, "mmaN needs to be multiple of 16 for 2CTA UTCMMA."); - } - - TLLM_CHECK_ERROR( - options.mTileM % options.mEpilogueTileM == 0 && options.mTileN % options.mEpilogueTileN == 0, - "TileM and TileN must be divisible by EpilogueTileM and EpilogueTileN respectively."); - TLLM_CHECK_ERROR( - (options.mClusterDimX == 1 || options.mClusterDimX == 2) && options.mClusterDimY == 1, - "GEMM does not support cluster in X and Y dimensions."); - TLLM_CHECK_ERROR(options.mClusterDimZ == 1 || options.mNumSlicesForSplitK > 1, - "Cluster DimZ is only allowed for split-k."); - TLLM_CHECK_ERROR(options.mTileM <= 128, "GEMM does not support TileM > 128."); - - // FIXME: this is a bug in DeepSeek Fp8. - if (options.mUseDeepSeekFp8) { - TLLM_CHECK_ERROR(options.mK % (options.mNumSlicesForSplitK * options.mTileK) == 0, - "K must be a multiple of TileK * numSlicesForSplitK for DeepSeekFp8"); - } - - // Check all-reduce options. - if (options.mAllReduceAlgo == AllReduceAlgo::OneShot) { - // One shot is implemented with PTX cp.reduce.async.bulk.tensor which supports only the - // following types for reduce add: u32, s32, u64, f32, f16, bf16. - // - // See: https://docs.nvidia.com/cuda/parallel-thread-execution/ - // #data-movement-and-conversion-instructions-cp-reduce-async-bulk-tensor - std::set dtypeSupported{tg::Dtype::UInt32, tg::Dtype::Int32, tg::Dtype::UInt64, - tg::Dtype::Fp32, tg::Dtype::Fp16, tg::Dtype::Bfloat16}; - TLLM_CHECK_ERROR(dtypeSupported.find(options.mDtypeC) != dtypeSupported.end(), - "Unsupported output dtype ", tg::dtypeToString(options.mDtypeC)); - } else if (options.mAllReduceAlgo == AllReduceAlgo::TwoShot) { - // TODO(anchengc): - // Input dtype == output dtype -> can perform all-reduce in-place. - // Input dtype != output dtype -> must perform all-reduce out of place. - TLLM_CHECK_ERROR_FMT( - options.mDtypeC == options.mDtypeAcc, - "Not implemented - mixed dtype (dtypeC (%s) != dtypeAcc (%s)) requires out of place update", - tg::dtypeToString(options.mDtypeC).c_str(), tg::dtypeToString(options.mDtypeAcc).c_str()); - } - if (options.mAllReduceAlgo != AllReduceAlgo::None) { - TLLM_CHECK_ERROR(options.mUseTmaStore, "Non-TMA store with all-reduce is not implemented"); - } - - if (updateOptions) { - if (options.mNumSlicesForSplitK == 1) { - // No split-k. - options.mSplitK = SplitK::None; - } else if (options.mNumSlicesForSplitK > 1 && options.mClusterDimZ == 1) { - // Split-k with exchange through gmem. - options.mSplitK = SplitK::Gmem; - } else { - // Split-k with exchange through Dsmem. - options.mSplitK = SplitK::Dsmem; - } - } - // For GMEM-based split-K, we write 4 elements at once. - if (options.mSplitK == SplitK::Gmem) { - TLLM_CHECK_ERROR((options.mM * options.mN) % 4 == 0, - "M * N must be a multiple of 4 for Split-K"); - } - - if (options.mNumSlicesForSplitK > 1) { - if ((options.mEpilogueTileM != options.mTileM || options.mEpilogueTileN != options.mTileN) && - !options.mUseDeepSeekFp8) { - TLLM_LOG_WARNING( - "Overwriting epilogueTileM and epilogueTileN to match tileM and tileN respectively"); - if (updateOptions) { - options.mEpilogueTileM = options.mTileM; - options.mEpilogueTileN = options.mTileN; - } else { - return false; - } - } - } - if (options.mSplitK == SplitK::Dsmem) { - TLLM_CHECK_ERROR(options.mClusterDimZ == options.mNumSlicesForSplitK, - "CGA size must be equal to the number of slices in split-k"); - } - - if (options.mUseShuffledMatrix && !options.mTransposeMmaOutput) { - TLLM_CHECK_ERROR( - !options.mUseDeepSeekFp8, - "DeepSeek Fp8 is not supported when using shuffled matrix and non-transposed mma output"); - TLLM_CHECK_ERROR( - options.mEpilogueLdtmBits == 32, - "EpilogueLdtmBits must be 32 when using shuffled matrix and non-transposed mma output"); - TLLM_CHECK_ERROR( - options.mEpilogueLdtmDps == 32, - "EpilogueLdtmDps must be 32 when using shuffled matrix and non-transposed mma output"); - TLLM_CHECK_ERROR( - options.mUseTmaStore, - "TMA store is required when using shuffled matrix and non-transposed mma output"); - TLLM_CHECK_ERROR( - !options.mSliceK, - "Slice-K is not supported when using shuffled matrix and non-transposed mma output"); - // When doing unshuffle in the epilogue, one fragment of epilogue tile must have at least one - // shuffle block. - auto minEpilogueTileN = getShuffleBlockSize(options.mEpilogueTileM); - TLLM_CHECK_ERROR_FMT(options.mEpilogueTileN >= minEpilogueTileN, - "EpilogueTileN (%d) must be a larger than the shuffle block size (%d) " - "when using shuffled matrix and non-transposed mma output", - options.mEpilogueTileN, minEpilogueTileN); - } - - // Maps numStagesMma to (stagesWithinWorkTile, stagesAcrossWorkTile) if not already set. - // If (-1, -1) -> (numStagesMma / min(2, numStagesMma), min(2, numStagesMma)) - // If ( m, -1) -> (m, numStagesMma / m) - // If (-1, n) -> (numStagesMma / n, n) - if (options.mNumStagesMmaWithinWorkTile == -1 && options.mNumStagesMmaAcrossWorkTile == -1) { - if (updateOptions) { - options.mNumStagesMmaAcrossWorkTile = std::min(2, options.mNumStagesMma); - options.mNumStagesMmaWithinWorkTile = - options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; - } else { - return false; - } - } else if (options.mNumStagesMmaWithinWorkTile == -1) { - if (updateOptions) { - options.mNumStagesMmaWithinWorkTile = - options.mNumStagesMma / options.mNumStagesMmaAcrossWorkTile; - } else { - return false; - } - } else if (options.mNumStagesMmaAcrossWorkTile == -1) { - if (updateOptions) { - options.mNumStagesMmaAcrossWorkTile = - options.mNumStagesMma / options.mNumStagesMmaWithinWorkTile; - } else { - return false; - } - } - // Check mma stages. - TLLM_CHECK_ERROR_FMT(options.mNumStagesMmaWithinWorkTile * options.mNumStagesMmaAcrossWorkTile == - options.mNumStagesMma && - options.mNumStagesMmaAcrossWorkTile <= 2, - "Condition numStagesMmaWithinWorkTile (%d) * numStagesMmaAcrossWorkTile " - "(%d) == numStagesMma (%d) && numStagesMmaAcrossWorkTile (%d) <= 2 must be " - "satisfied. Check arguments.", - options.mNumStagesMmaWithinWorkTile, options.mNumStagesMmaAcrossWorkTile, - options.mNumStagesMma, options.mNumStagesMmaAcrossWorkTile); - // Mma stage must be 1 for pre-Hopper. - TLLM_CHECK_ERROR(isBlackwell || options.mNumStagesMma == 1, - "Mma stage must be 1 for pre-Hopper. Found ", options.mNumStagesMma); - // DeepSeek Fp8 - if (!options.mUseDeepSeekFp8) { - TLLM_CHECK_ERROR(options.mNumStagesMmaWithinWorkTile == 1, - "Non-DeepSeekFp8 requires numStagesMmaWithinWorkTile == 1"); - if (options.mNumStagesMma > 1) { - TLLM_CHECK_ERROR(options.mTileScheduler == TileScheduler::Persistent, - "Non-DeepSeekFp8 requires persistent scheduler when using numStagesMma >1"); - } - } - if (options.mUseDeepSeekFp8) { - TLLM_CHECK_ERROR(options.mClusterDimX == 1, "2CTA Gemm is not supported for DeepSeekFp8"); - } - if (options.mUseDeepSeekFp8) { - TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, - "A and B dtype must be E4m3 for DeepSeek Fp8. Found dtypeA=", - tg::dtypeToString(options.mDtypeA), - " dtypeB=", tg::dtypeToString(options.mDtypeB)); - - TLLM_CHECK_ERROR(isBlackwell, "DeepSeek Fp8 is not supported for Hopper"); - TLLM_CHECK_ERROR(options.mAllReduceAlgo == AllReduceAlgo::None, - "DeepSeek Fp8 does not support AllReduce"); - - // Check that TileK = 128 for correct scaling of every 128 channels. - TLLM_CHECK_ERROR(options.mTileK == 128, "Tile-K must be equal to 128 for DeepSeek Fp8"); - TLLM_CHECK_ERROR(options.mK % options.mTileK == 0, "K must be a multiple of TileK"); - // Tile sizes of the output hidden dimension. - auto hiddenDimPerOutputTile = options.mTransposeMmaOutput ? options.mTileM : options.mTileN; - auto hiddenDimPerEpilogueTile = - options.mTransposeMmaOutput ? options.mEpilogueTileM : options.mEpilogueTileN; - auto hiddenDimPerMma = options.mTransposeMmaOutput ? options.mMmaM : options.mMmaN; - auto hiddenDimName = options.mTransposeMmaOutput ? "M" : "N"; - TLLM_CHECK_WARNING(options.mNumStagesMmaWithinWorkTile > 1, - "DeepSeekFp8 recommends setting \"-numStagesMmaWithinWorkTile 2\"."); - // Update the number of stages of the MMA accumulator pipeline. TODO: enable by default for - // deepseek. - // options.mNumStagesMma = 2; - // Use two MMA warps to reduce mbar trywait latency. TODO: enable by default for deepseek. - // options.mUseTwoMmaWarps = true; - - // Make sure the GEMM-K dimension is a multiple of 128 when using DeepSeek FP8. - TLLM_CHECK_ERROR(options.mK % 128 == 0 && options.mValidK % 128 == 0, - "GEMM-K and validK must be a multiple of 128 when using DeepSeek Fp8. Found ", - options.mK, " and validK=", options.mValidK); - - // Check that the output tile N can be processed with the epilogue tile granularity. - TLLM_CHECK_ERROR((hiddenDimPerOutputTile / 2) % hiddenDimPerEpilogueTile == 0, - "DeepSeek Fp8 requires Tile", hiddenDimName, " / 2 (", - hiddenDimPerOutputTile / 2, ") being a multiple of EpilogueTile", - hiddenDimName, " (", hiddenDimPerEpilogueTile, ")"); - // Check that the output tile N can be processed with the epilogue tile granularity. - TLLM_CHECK_ERROR((hiddenDimPerOutputTile / 2) % hiddenDimPerMma == 0, - "DeepSeek Fp8 requires Tile", hiddenDimName, " / 2 (", - hiddenDimPerOutputTile / 2, ") being a multiple of mma", hiddenDimName, " (", - hiddenDimPerMma, ")"); - } - - TLLM_CHECK_ERROR(options.mNumEpilogueWarps == 4 || options.mNumEpilogueWarps == 8, - "mNumEpilogueWarps has to be either 4 or 8."); - - if (options.mSliceK) { - TLLM_CHECK_ERROR(isBlackwell, "Slice-K is not supported on Hopper"); - - TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, "DeepSeek Fp8 GEMM is not supported for slice-K"); - TLLM_CHECK_ERROR(options.mUseTwoTmaLoadWarps, "Slice-K requires two warp load for A and B"); - TLLM_CHECK_ERROR(options.mTransposeMmaOutput, "Slice-K requires transpose mma output"); - TLLM_CHECK_ERROR(options.mUseShuffledMatrix, "Slice-K requires shuffled matrix"); - TLLM_CHECK_ERROR(options.mTileK % 128 == 0, "Slice-K requires TileK be a multiple of 128"); - TLLM_CHECK_ERROR(options.mMmaM == 128, "Slice-K requires MmaM == 128"); - TLLM_CHECK_ERROR(options.mTileN == options.mEpilogueTileN, - "TileN must be equal to EpilogueTileN for slice-K"); - - TLLM_LOG_WARNING("Overwriting TileM and EpilogueTileM to 32 for slice-K"); - if (options.mTileM != 32 || options.mEpilogueTileM != 32) { - if (updateOptions) { - // FIXME: it is possible to remove this restriction. - options.mTileM = 32; - options.mEpilogueTileM = 32; - } else { - return false; - } - } - TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, - "Slice-K requires e4m3 input dtype"); - - if (options.mNumSlicesForSliceK != 4) { - if (updateOptions) { - options.mNumSlicesForSliceK = 4; - } else { - return false; - } - } - TLLM_CHECK_ERROR((options.mTileK / options.mMmaK) % options.mNumSlicesForSliceK == 0, "TileK (", - options.mTileK, ") / MmaK (", options.mMmaK, - ") must be a multiple of mNumSlicesForSliceK (", options.mNumSlicesForSliceK, - ")"); - } - - // Number of iterations in K dimension after padding. - // Note the perCtaK in each CTA in the splitK group are padded to the same number of iterations. - // E.g., K = 512, TileK = 128, numSlicesForSplitK = 3. Then the padded K is - // - // ceil(512 / (128*3)) * (128*3) = 768 - // - int const paddedK = divUpMul(options.mK, options.mTileK * options.mNumSlicesForSplitK); - int const perCtaK = paddedK / options.mNumSlicesForSplitK; - // However, number of iterations is clamped to multiples of tileK within individual CTAs - // E.g., K = 448, TileK = 64, numSlicesForSplitK = 4. - // - // paddedK = 512 - // perCtaK = 128 - // clampedPerCtaK for CTA 0, 1, 2 = 128 - // clampedPerCtaK for CTA 3 = 64 - int const paddingForK = paddedK - options.mK; - int const clampedAndPaddedPerCtaK = divUpMul(perCtaK - paddingForK, options.mTileK); - if (options.mUseUnrollLoop2xForMma) { - // Check that the padded K and clamped padded K (K rounded to next multiple of tileK) is a - // multiple of 2*TileK when UnrollLoop2x is enabled. This is to avoid deadlock when mma runs - // even-numbered loop while the other warps run odd-numbered loop. - // - bool notSupported = (perCtaK % (options.mTileK * 2) != 0) || - (clampedAndPaddedPerCtaK % (options.mTileK * 2) != 0); - if (notSupported) { - TLLM_LOG_WARNING("Size K / splitK must be a multiple of TileK * 2. Found TileK=", - options.mTileK, " and K=", options.mK, " (paddedK=", paddedK, - " clampedAndPaddedPerCtaK=", clampedAndPaddedPerCtaK, - ") and numSlicesForSplitK=", options.mNumSlicesForSplitK, - ". Disabling unrollLoop2xForMma."); - if (updateOptions) { - options.mUseUnrollLoop2xForMma = false; - } else { - return false; - } - } - } - if (options.mNumSlicesForSplitK > 1) { - TLLM_CHECK_ERROR( - perCtaK * (options.mNumSlicesForSplitK - 1) < options.mK, - "K must be greater than perCtaK * (numSlicesForSplitK - 1) to ensure each CTA has work"); - } - - if (!isBlackwell && options.mTileScheduler == TileScheduler::Persistent) { - // TODO(anchengc): will be supported in upcoming MRs. - TLLM_LOG_WARNING("Persistent scheduling is not supported on Hopper. Using Static scheduling."); - if (updateOptions) { - options.mTileScheduler = TileScheduler::Static; - } else { - return false; - } - } - - if (isBlackwell && !options.mUseCustomMmaSchedule && !options.mUseDeepSeekFp8 && - options.mTileScheduler == TileScheduler::Persistent) { - if (updateOptions) { - options.mUseCustomMmaSchedule = true; - } else { - TLLM_CHECK_ERROR(false, - "TileScheduler::Persistent and !UseCustomMmaSchedule is not supported."); - } - } - - if (options.mEnablesDelayedEarlyExit && options.mEnablesEarlyExit) { - TLLM_LOG_WARNING( - "Only one of early exit and delayed early exit should be enabled. Disabling " - "delayed early exit"); - if (updateOptions) { - options.mEnablesDelayedEarlyExit = false; - } else { - return false; - } - } - - // This check prevents the triggering of the secondary (PREEXIT) from executing before the wait - // for primary (ACQBULK). This could lead to the following confusing situation, which we want to - // avoid: - // - // Kernel 3 is written with the assumption that it can read the output of - // kernel 1 *without* ACQBULK and the output of kernel 2 *with* ACQBULK. - // However, when we allow PREEXIT and ACQBULK to be executed out of order, - // this is not guaranteed. - // - // Time: ----> - // - // Kernel 1: ----PREEXIT-----------FLUSH - // Kernel 2: -------PREEXIT----ACQBULK---FLUSH - // Kernel 3: Warp 0: ---- (!) Output of 1,2 is not yet visible - // ----------------------- - // Warp 1: ---- (!) We normally assume that 1 is visible is not yet - // visible- Warp 2: -------------------ACQBULK-- Kernel 1,2 output visible - // ---------- - TLLM_CHECK_ERROR( - (options.mGridWaitForPrimaryA || !options.mGridTriggerSecondaryA), - "A: If a task triggers a secondary kernel, it must also wait for primary kernel."); - TLLM_CHECK_ERROR( - (options.mGridWaitForPrimaryB || !options.mGridTriggerSecondaryB), - "B: If a task triggers a secondary kernel, it must also wait for primary kernel."); - - if (options.mUsePerTokenSfA || options.mUsePerTokenSfB) { - // Checks applicable to both MetaFP8 and RoutingScalesOnInput - TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, - "DeepSeek FP8 and per-token scaling are not compatible"); - TLLM_CHECK_ERROR(isBlackwell, "Per-token scaling is not supported for Hopper"); - if (options.mUsePerTokenSfA && options.mUsePerTokenSfB) { - // MetaFP8 case - TLLM_CHECK_ERROR(options.mDtypeA == tg::Dtype::E4m3 && options.mDtypeB == tg::Dtype::E4m3, - "A and B dtype must be E4m3 for Meta Fp8. Found dtypeA=", - tg::dtypeToString(options.mDtypeA), - " dtypeB=", tg::dtypeToString(options.mDtypeB)); - } else { - // RoutingScalesOnInput case - TLLM_CHECK_ERROR((options.mUsePerTokenSfA && !options.mTransposeMmaOutput) || - (options.mUsePerTokenSfB && options.mTransposeMmaOutput), - "In RoutingScalesOnInput mode, perToken scales must be used on activations"); - } - } - - // The generation should support non K-major layouts for both A and B; however, it is unclear if - // there is a use-case - TLLM_CHECK_ERROR( - (options.mLayoutA == MatrixLayout::MajorK) || (options.mLayoutB == MatrixLayout::MajorK), - "At least one matrix must be in k-major layout"); - - // Some features are currently only support when both matrices are in K-major format - if (options.mLayoutA != MatrixLayout::MajorK || options.mLayoutB != MatrixLayout::MajorK) { - TLLM_CHECK_ERROR(isBlackwell, "Non K-major layouts are only supported on Blackwell"); - TLLM_CHECK_ERROR(options.mSplitK == SplitK::None, "Non K-major layouts do not support split K"); - } - if (options.mLayoutA == MatrixLayout::MajorMn) { - TLLM_CHECK_ERROR(tg::dtypeGetNumBits(options.mDtypeA) >= 8, - "Subbyte types do not support m-major layout"); - } - if (options.mLayoutB == MatrixLayout::MajorMn) { - TLLM_CHECK_ERROR(tg::dtypeGetNumBits(options.mDtypeB) >= 8, - "Subbyte types do not support n-major layout"); - } - - if ((options.mLayoutA == MatrixLayout::BlockMajorK) || - (options.mLayoutB == MatrixLayout::BlockMajorK)) { - bool const isBlockA = options.mLayoutA == MatrixLayout::BlockMajorK; - - // Block K size must be 128B. - // TODO Leaving this as an option for now in case we want to expertiment with other block sizes - // As the user is not expected to set this, do not fail if updateOptions is false - int32_t const elemSizeInBits = - (isBlockA) ? tg::dtypeGetNumBits(options.mDtypeA) : tg::dtypeGetNumBits(options.mDtypeB); - int32_t const elemsIn128B = 128 * 8 /* Bits in byte */ / elemSizeInBits; - - if (options.mBlockK != elemsIn128B) { - if (updateOptions) { - options.mBlockK = elemsIn128B; - } else { - return false; - } - } - - if (options.mBlockK > options.mTileK) { - TLLM_CHECK_ERROR( - options.mBlockK % options.mTileK == 0, - "If block size is greater than tile size, block size must be a multiple of tile size"); - } else if (options.mBlockK < options.mTileK) { - TLLM_CHECK_ERROR( - options.mTileK % options.mBlockK == 0, - "If tile size is greater than block size, tile size must be a multiple of block size"); - } - } - - if (!isBiasTypeNone(options.mBiasType)) { - TLLM_CHECK_ERROR(!isBiasTypeMn(options.mBiasType), "BiasType::Mn is not supported"); - TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, "Bias is not supported for DeepSeek Fp8"); - TLLM_CHECK_ERROR(!(options.mUsePerTokenSfA && options.mUsePerTokenSfB), - "Bias is not supported for Meta Fp8"); - } - - if (options.mUseMaxTmemOverlap) { - TLLM_CHECK_ERROR(options.mUseTmaStore, "mUseMaxTmemOverlap only works with TMA store"); - TLLM_CHECK_ERROR(options.mFuseUtccpWithUtcmma, - "mUseMaxTmemOverlap only works with mFuseUtccpWithUtcmma"); - TLLM_CHECK_ERROR(options.mNumSlicesForSplitK == 1, - "mUseMaxTmemOverlap does not work with splitK"); - TLLM_CHECK_ERROR(options.mNumSlicesForSliceK == 1, - "mUseMaxTmemOverlap does not work with sliceK"); - TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, - "mUseMaxTmemOverlap does not work with mUseDeepSeekFp8"); - TLLM_CHECK_ERROR(!options.mUseUnrollLoop2xForMma, - "mUseMaxTmemOverlap does not work with mUseUnrollLoop2xForMma"); - } - - if (options.mNumEpilogueWarps > 4) { - TLLM_CHECK_ERROR(options.mUseTmaStore, - "Using more than 4 warps for epilogue only works with TMA store"); - TLLM_CHECK_ERROR(options.mNumSlicesForSplitK == 1, - "Using more than 4 warps for epilogue does not work with splitK"); - TLLM_CHECK_ERROR(options.mNumSlicesForSliceK == 1, - "Using more than 4 warps for epilogue does not work with sliceK"); - TLLM_CHECK_ERROR(!options.mUseDeepSeekFp8, - "Using more than 4 warps for epilogue does not work with mUseDeepSeekFp8"); - - auto const numEpilogueWrpGrps = options.mNumEpilogueWarps / 4; - TLLM_CHECK_ERROR(options.mTileN % (options.mEpilogueTileN * numEpilogueWrpGrps) == 0, - "TileN must be a multiple of EpilogueTileN * numEpilogueWrpGrps"); - } - - if (updateOptions) { - // Init kernel traits. - options.mKernelTraits = KernelTraits( - options.mDtypeA, options.mDtypeB, options.mDtypeC, options.mDtypeAcc, options.mDtypeMmaA, - options.mDtypeMmaB, options.mMmaKind, options.mSparsityA, options.mMmaK, options.mTileM, - options.mTileN, options.mTileK, options.mEpilogueTileM, options.mEpilogueTileN, - options.mSfBlockSizeA, options.mSfBlockSizeB, options.mNumStages, options.mNumStagesMma, - options.mNumSlicesForSplitK, options.mNumSlicesForSliceK, options.mSplitK, - options.mUseTmaStore, options.mTransposeMmaOutput, options.mAllReduceAlgo, - options.mFuseUtccpWithUtcmma, options.mUseMaxTmemOverlap, options.mNumEpilogueWarps, - options.mTileScheduler == TileScheduler::Persistent, options.mUseDeepSeekFp8, - options.mUsePerTokenSfA, options.mUsePerTokenSfB, - /* useTwoCtas*/ options.mClusterDimX == 2, options.mBiasType); - } - - return true; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool getDoesScaleC(tg::Dtype dtypeC) { - // Need to scale/quantize the output C matrix when the output type is Fp8 or NvFp4. - return dtypeC == tg::Dtype::E4m3 || dtypeC == tg::Dtype::E2m1; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool getDoesScaleAb(tg::Dtype dtypeA, tg::Dtype dtypeB, bool useDeepSeekFp8) { - // Need to scale/dequantize the input A/B matrices when the input type is Fp8 or NvFp4 and - // DeepSeekFp8 is not used. - bool const doesScaleAb{ - dtypeA == tg::Dtype::E2m1 || dtypeB == tg::Dtype::E2m1 || - ((dtypeA == tg::Dtype::E4m3 || dtypeB == tg::Dtype::E4m3) && !useDeepSeekFp8)}; - return doesScaleAb; -} - -////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool getDoesScaleAct(tg::Dtype dtypeA, tg::Dtype dtypeB, bool useDeepSeekFp8, - EltwiseActType eltwiseActType) { - // Only non-linear activations require separate scaleAct. - bool const isLinearAct = eltwiseActType == EltwiseActType::None; - return !isLinearAct && getDoesScaleAb(dtypeA, dtypeB, useDeepSeekFp8); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool getKernelDoesScaleC(tg::Dtype dtypeA, tg::Dtype dtypeB, tg::Dtype dtypeC, - bool useDeepSeekFp8) { - // In the Gemm/BatchedGemm kernels, dequantScaleAb and quantScaleC are combined into one single - // scaling factor (called scaleC). As a result, we combine the logic for getDoesScaleAb and - // getDoesScaleC. - return getDoesScaleC(dtypeC) || getDoesScaleAb(dtypeA, dtypeB, useDeepSeekFp8); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline CUresult loadCubinData(CUmodule* module, Config const& config) { - // Trtllm links the cubin into the executable while Flashinfer loads the cubin from storage. -#ifdef TLLM_GEN_EXPORT_FLASHINFER -#ifdef TLLM_GEN_GEMM_CUBIN_PATH - static const std::string tllm_gen_gemm_cubin_path = std::string(TLLM_GEN_GEMM_CUBIN_PATH); - const std::string sha256 = config.mHash ? config.mHash : ""; - std::string fileName = config.mFunctionName; - if (!fileName.empty()) { - fileName[0] = static_cast(std::toupper(static_cast(fileName[0]))); - } - const std::string& data = flashinfer::trtllm_cubin_loader::getCubin( - tllm_gen_gemm_cubin_path + "/" + fileName + ".cubin", sha256); - CUresult result = cuModuleLoadData(module, data.c_str()); -#else - static_assert(false, "TLLM_GEN_GEMM_CUBIN_PATH macro is not defined when compiling"); -#endif // TLLM_GEN_GEMM_CUBIN_PATH -#else - CUresult result = cuModuleLoadData(module, config.mData); -#endif // TLLM_GEN_EXPORT_FLASHINFER - return result; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gemm - -#ifdef TLLM_GEN_EXPORT_INTERFACE - -#undef TLLM_CHECK_ERROR -#undef TLLM_CHECK_ERROR_FMT -#undef TLLM_CHECK_WARNING -#undef TLLM_LOG_WARNING -#undef TLLM_LOG_INFO -#undef TLLM_LOG_ERROR - -#endif // TLLM_GEN_EXPORT_INTERFACE - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h deleted file mode 100644 index bda17264fb..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParams.h +++ /dev/null @@ -1,679 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include - -#include "BatchedGemmEnums.h" -#include "Enums.h" -#include "TmaDescriptor.h" -#include "trtllm/gen/CommonUtils.h" -#include "trtllm/gen/SfLayoutDecl.h" -#include "trtllm/gen/SparsityDecl.h" - -// NOTE: keep this code dependency free. It has to be included by the device code and has to be -// compilable with NVRTC. -#include "KernelParamsDecl.h" - -namespace batchedGemm { - -namespace batchedGemm { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// TODO: Find a better header to put this in, that we can include from here. -template -inline T ceilDiv(T m, T n) { - return (m + n - T(1)) / n; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace tg = trtllm::gen; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace KernelParamsSetup { -#ifdef TLLM_ENABLE_CUDA - -enum class MatrixType { MatrixA = 0, MatrixB, MatrixC }; - -////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Utility functions. -// -////////////////////////////////////////////////////////////////////////////////////////////////// - -template -bool useTmaOobOptA(BatchedGemmOptions const& options) { - return options.mBatchMode == BatchedGemmOptions::BatchMode::BatchM && - doesRouteImplUseNoRoute(options.mRouteImpl) && options.mUseTmaOobOpt; -} - -////////////////////////////////////////////////////////////////////////////////////////////////// - -template -bool useTmaOobOptB(BatchedGemmOptions const& options) { - return options.mBatchMode == BatchedGemmOptions::BatchMode::BatchN && - doesRouteImplUseNoRoute(options.mRouteImpl) && options.mUseTmaOobOpt; -} - -////////////////////////////////////////////////////////////////////////////////////////////////// - -template -bool useTmaOobOptC(BatchedGemmOptions const& options) { - return options.mUseTmaStore && options.mUseTmaOobOpt; -} - -////////////////////////////////////////////////////////////////////////////////////////////////// - -// Create the TMA shape/stride for A/B/C. -template -static auto makeTmaShapeStrideAbc(GemmOptions const& options, int sizeM, int sizeN, int sizeK, - int tileM, int tileN, int tileK, MatrixType matrixType, - int validM = -1, int validN = -1, int validK = -1) { - // Default to padded dimensions if not provided. - validM = validM < 0 ? sizeM : validM; - validN = validN < 0 ? sizeN : validN; - validK = validK < 0 ? sizeK : validK; - // Weights matrix is A if we transpose the output of MMA (to have it M-major). - // Otherwise, it is B, when the output of MMA is K-major. - bool const isWeights = (matrixType == MatrixType::MatrixA && options.mTransposeMmaOutput) || - (matrixType == MatrixType::MatrixB && !options.mTransposeMmaOutput); - - // Whether to use TMA OOB trick to block out padded dummy tokens and saving BW whenever no routing - // is involved. It applies to batchM and matrixA, or batchN and matrixB, or any case for matrixC. - bool const useTmaOobOpt = matrixType == MatrixType::MatrixA ? useTmaOobOptA(options) - : matrixType == MatrixType::MatrixB ? useTmaOobOptB(options) - : matrixType == MatrixType::MatrixC ? useTmaOobOptC(options) - : false; - - // The outer dimension. Uses padded dimensions for strides and valid dimensions for shapes. - auto numTokens = - (matrixType == MatrixType::MatrixA || matrixType == MatrixType::MatrixC) ? sizeM : sizeN; - auto numTokensValid = - (matrixType == MatrixType::MatrixA || matrixType == MatrixType::MatrixC) ? validM : validN; - // The outer dimension tile size. - auto ctaTileNumTokens = - (matrixType == MatrixType::MatrixA || matrixType == MatrixType::MatrixC) ? tileM : tileN; - // The outer dimension of TMA box shape. - auto tileNumTokens = - (matrixType == MatrixType::MatrixC) ? options.mEpilogueTileM : ctaTileNumTokens; - - // The inner dimension. - auto hiddenSize = (matrixType == MatrixType::MatrixC) ? sizeN : sizeK; - auto hiddenSizeValid = (matrixType == MatrixType::MatrixC) ? validN : validK; - // The inner dimension tile size. - auto ctaTileHiddenSize = (matrixType == MatrixType::MatrixC) ? tileN : tileK; - // The inner dimension of TMA box shape. - auto tileHiddenSize = - (matrixType == MatrixType::MatrixC) ? options.mEpilogueTileN : ctaTileHiddenSize; - - // Swap matrix C sizes if output is transposed. - if (matrixType == MatrixType::MatrixC && options.mTransposeMmaOutput) { - std::swap(numTokens, hiddenSize); - std::swap(numTokensValid, hiddenSizeValid); - std::swap(ctaTileNumTokens, ctaTileHiddenSize); - std::swap(tileNumTokens, tileHiddenSize); - } - - // For a fused activation kernel, the hidden size of output is halved. TODO: That's true for - // gated activations but not regular activations. - if (options.mFusedAct && matrixType == MatrixType::MatrixC) { - hiddenSize /= 2; - hiddenSizeValid /= 2; - tileHiddenSize /= 2; - ctaTileHiddenSize /= 2; - } - - // The cute tensor shape for A/B: (numTokens, hiddenSize). - // Note that TMA descriptor expects the first dimension's stride to be - // 1, so swap the first two dimension so that the hiddenSize dimension comes first. - - // Activations matrix is 2D (sum(divUpMul(M[bi], tileM) for bi in B), K). - // Use valid dimensions for shape. - std::vector shape = {static_cast(hiddenSizeValid), - static_cast(numTokensValid)}; - if (useTmaOobOpt /* also implies input/output activation */) { - // If TMA OOB optimization is used: - // Shape [hidden, tokens] Stride [1, hidden] becomes - // Shape [hidden, tileN, TmaDimMax, TmaDimMax] Stride [1, hidden, XLargeN - hidden, hidden] - shape = {static_cast(hiddenSizeValid), static_cast(ctaTileNumTokens), - static_cast(tg::TmaDimMax), static_cast(tg::TmaDimMax)}; - } else if (isWeights) { - // If the matrix is a weights matrix, we use 3D logical shape (B, M, K) or (B, N, K). - shape = {static_cast(hiddenSizeValid), static_cast(numTokensValid), - static_cast(options.mNumBatches)}; - } - - // Assemble the stride (strideTokens, 1). - // Swap the first two dimension as mentioned before. - std::vector stride = {1, static_cast(hiddenSize)}; - if (useTmaOobOpt) { - stride = {1, static_cast(hiddenSize), static_cast(tg::XLargeN - hiddenSize), - static_cast(hiddenSize)}; - } else if (isWeights) { - stride = {1, static_cast(hiddenSize), - static_cast(hiddenSize) * static_cast(numTokens)}; - } - - // Assemble the box shape - std::vector tileShape = {tileHiddenSize, tileNumTokens}; - - // Alternate layouts (MajorMn and BlockMajorK) do not apply to matrixC - if (matrixType != MatrixType::MatrixC) { - // When using 2CTA MMA, we only need to load half of the tile in each CTA for B. - if (matrixType == MatrixType::MatrixB && tileShape[1] > 1 && options.mClusterDimX == 2) { - tileShape[1] /= 2; - } - gemm::MatrixLayout layout = - (matrixType == MatrixType::MatrixA) ? options.mLayoutA : options.mLayoutB; - // Note, only the weights support non MajorK layouts - if (layout == gemm::MatrixLayout::MajorMn) { - // Apply transpose if necessary - std::swap(shape[0], shape[1]); - stride[1] = numTokens; - std::swap(tileShape[0], tileShape[1]); - } else if (layout == gemm::MatrixLayout::BlockMajorK) { - // Set shapes based on blocking layout. - shape = {static_cast(options.mBlockK), static_cast(numTokens), - static_cast(sizeK / options.mBlockK), - static_cast(options.mNumBatches)}; - // Strides use padded dimensions - stride = {1, static_cast(options.mBlockK), - static_cast(numTokens * options.mBlockK), - static_cast(hiddenSize * numTokens)}; - - // If blockK > tileK, then the inner most box size will be based on the tile - int32_t const tileBlockK = std::min(options.mBlockK, tileHiddenSize); - tileShape = {tileBlockK, tileNumTokens, tileHiddenSize / tileBlockK}; - } - } - - return std::make_tuple(shape, stride, tileShape); -} - -// Create the TMA shape/stride for A/B block scaling factors. -static auto makeTmaShapeStrideSfAb(int mM, int mN, int mK, MatrixType matrixType, int tileM, - int tileN, int tileK, tg::SfLayout layout, int sfReshapeFactor, - const int32_t numEltsPerSf) { - // The outer dimension. - auto numTokens = matrixType == MatrixType::MatrixA ? mM : mN; - // The inner dimension. - auto hiddenSize = mK; - // The outer tile dimension. - auto numTokensPerTile = matrixType == MatrixType::MatrixA ? tileM : tileN; - // The inner tile dimension. - auto hiddenSizePerTile = tileK; - - switch (layout) { - case tg::SfLayout::R128c4: { - auto shape = std::vector{ - 256, 2, static_cast(ceilDiv(hiddenSize, numEltsPerSf * 4)), - static_cast(ceilDiv(numTokens, 128))}; - - std::vector stride(shape.size()); - stride[0] = 1; - for (size_t i = 1; i < shape.size(); i++) { - stride[i] = shape[i - 1] * stride[i - 1]; - } - - auto tileShapes = std::vector{ - 256, 2, static_cast(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4)), - static_cast(ceilDiv(numTokensPerTile, 128))}; - - return std::make_tuple(shape, stride, tileShapes); - } - - case tg::SfLayout::R8c4: { - // The scaling factor tensor packs 8x4 tiles into contiguous 32B blocks. - // - // As the inner dimension (k) is often a multiple of the tile size, we can reshape to use - // fewer read requests, if the tile dimensions allow. It does not reduce the number of - // instructions. - // - // I.e., let's define r = min(⌈hiddenSizePerTile / (numEltsPerSf * 4)⌉, 8) - // - // The "logical" tensor is: [outer, inner / numEltsPerSf] - // The 8x4 SF layout is: [⌈outer / 8⌉, inner / (4 * numEltsPerSf), 32] - // The TMA tensor shape is: [⌈outer / 8⌉, inner / (4 * numEltsPerSf * r), r * 32] - // - // The caveat of NumRepeats>1 is we must pad the hidden dimension of SF to multiples of - // NumRepeats * numEltsPerSf * 4. - - // Detect if the supplied factor is power of 2. E.g., 0b0100 and (0b0100 - 1) == 0b0000. - int const r = sfReshapeFactor; - if (r > 0 && (r & (r - 1)) != 0) { - throw std::runtime_error("mSfReshapeFactor must be positive and a power of 2. Found " + - std::to_string(r)); - } - - // Sanitize number of repeats so it doesn't exceed the dimension. - int const repeats = std::min(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4), r); - - // Detect if the input hidden size K is a multiple of the repeats. - if (ceilDiv(hiddenSize, numEltsPerSf * 4) % repeats != 0) { - throw std::runtime_error( - "SF hiddenSize K (" + std::to_string(ceilDiv(hiddenSize, numEltsPerSf * 4)) + - ") must be a multiple of repeats (" + std::to_string(repeats) + ")"); - } - - auto shape = std::vector{ - static_cast(repeats * 32), - static_cast(ceilDiv(hiddenSize, numEltsPerSf * 4 * repeats)), - static_cast(ceilDiv(numTokens, 8))}; - - std::vector stride(shape.size()); - stride[0] = 1; - for (size_t i = 1; i < shape.size(); i++) { - stride[i] = shape[i - 1] * stride[i - 1]; - } - - auto tileShapes = std::vector{ - static_cast(repeats * 32), - static_cast(ceilDiv(hiddenSizePerTile, numEltsPerSf * 4 * repeats)), - static_cast(ceilDiv(numTokensPerTile, 8))}; - - return std::make_tuple(shape, stride, tileShapes); - } - - default: - throw std::runtime_error("Unsupported SF layout"); - } - return std::make_tuple(std::vector{}, std::vector{}, std::vector{}); -} - -// Create the TMA shape/stride for sparsity information of A. -template -static auto makeTmaShapeStrideSparsityInfoA(GemmOptions const& options) { - // Tensor dimensions. - auto outerDim = options.mM; - auto innerDim = tg::getNumBytesSparsityInfo(options.mSparsityA, options.mK); - // Tile dimensions. - auto tileOuterDim = options.mTileM; - auto tileInnerDim = tg::getNumBytesSparsityInfo(options.mSparsityA, options.mTileK); - - auto shape = - std::vector{static_cast(innerDim), static_cast(outerDim), - static_cast(options.mNumBatches)}; - - std::vector stride(shape.size()); - stride[0] = 1; - for (size_t i = 1; i < shape.size(); i++) { - stride[i] = shape[i - 1] * stride[i - 1]; - } - - auto tileShapes = - std::vector{static_cast(tileInnerDim), static_cast(tileOuterDim)}; - - return std::make_tuple(shape, stride, tileShapes); -} - -template -static KernelParams setKernelParams( - GemmOptions_ const& options, bool const batchM, void const* ptrA, void const* ptrB, void* ptrC, - void const* dSfA, void const* dSfB, void const* ptrPerTokenSfA, void const* ptrPerTokenSfB, - [[maybe_unused]] void const* ptrSparsityInfoA, void const* ptrBias, void* dSfC, - float const* ptrScaleC, float const* ptrScaleAct, float const* ptrScaleGate, - float const* ptrClampLimit, float const* ptrGatedActAlpha, float const* ptrGatedActBeta, - int32_t const* routeMap, float* rowMax, uint32_t* rowMaxBars, - int32_t const* ptrNumNonExitingCtas = nullptr, int32_t const* ptrTotalNumPaddedTokens = nullptr, - int32_t const* ptrCtaIdxXyToBatchIdx = nullptr, int32_t const* ptrCtaIdxXyToMnLimit = nullptr, - int32_t const maxNumCtas = KernelParams::MaxNumCtas) { - static_assert(sizeof(KernelParams) <= 32 * 1024, - "sizeof(KernelParams) has to be less or equal than 32KB"); - - // Create the return struct. - KernelParams params; - - params.ptrRouteMap = routeMap; - params.numTokens = options.mNumTokens; - - params.ptrScaleC = ptrScaleC; - params.ptrScaleAct = ptrScaleAct; - params.ptrScaleGate = ptrScaleGate; - params.ptrClampLimit = ptrClampLimit; - params.ptrGatedActAlpha = ptrGatedActAlpha; - params.ptrGatedActBeta = ptrGatedActBeta; - - int32_t ctaOffset = 0; - - int32_t const isSparseA = tg::isSparse(options.mSparsityA); - - // Compute totalNumPaddedTokens, ctaIdxXyToBatchIdx and ctaIdxXyToMnLimit if the batch dims are - // known at kernel launch time. Otherwise, these parameters are defined in the device buffers: - // ptrTotalNumPaddedTokens, ptrCtaIdxXyToBatchIdx and ptrCtaIdxXyToMnLimit respectively. - - if (options.mIsStaticBatch && !options.mIsUniformNumTokensPerBatch) { - params.totalNumPaddedTokens = 0; - for (int b = 0; b < options.mNumBatches; b++) { - int mM = batchM ? options.mBatchedM[b] : options.mM; - int mN = batchM ? options.mN : options.mBatchedN[b]; - - // Skip Tma descriptor creation if expert isn't used - if (mM == 0 || mN == 0) { - continue; - } - - // The number of CTAs. - int32_t numCtas = batchM ? (mM + options.mTileM * options.mClusterDimX - 1) / - (options.mTileM * options.mClusterDimX) * options.mClusterDimX - : (mN + options.mTileN - 1) / options.mTileN; - // The size of the tile. - int32_t tile = batchM ? options.mTileM * options.mClusterDimX : options.mTileN; - // The problem size. - int32_t mn = batchM ? mM : mN; - int32_t tokensPerTile = mn; - - // Make sure we do not exceed the launch limit. - if (ctaOffset + numCtas > KernelParams::MaxNumCtas) { - throw std::runtime_error("Too many CTAs"); - } - - for (int32_t cta = 0; cta < numCtas; cta++) { - params.ctaIdxXyToBatchIdx[ctaOffset + cta] = b; - // This is now an identity map and it is no longer needed. - // params.ctaIdxXyToTileIdxMn[ctaOffset + cta] = ctaOffset + cta; - params.ctaIdxXyToMnLimit[ctaOffset + cta] = - std::min((ctaOffset + cta + 1) * tile, ctaOffset * tile + tokensPerTile); - } - ctaOffset += numCtas; - - params.totalNumPaddedTokens += numCtas * tile; - } - params.totalNumOutputPaddedTokens = params.totalNumPaddedTokens; - } else if (options.mIsStaticBatch && options.mIsUniformNumTokensPerBatch) { - auto numTokens = batchM ? options.mBatchedM[0] : options.mBatchedN[0]; - auto tileTokensDim = batchM ? options.mTileM * options.mClusterDimX : options.mTileN; - params.batchStrideInCtas = (options.mBatchStrideInTokens + tileTokensDim - 1) / tileTokensDim; - params.ctasInTokenDimPerBatch = (numTokens + tileTokensDim - 1) / tileTokensDim; - params.totalNumOutputPaddedTokens = - params.ctasInTokenDimPerBatch * tileTokensDim * options.mNumBatches; - if (params.batchStrideInCtas == 0) { - params.totalNumPaddedTokens = params.ctasInTokenDimPerBatch * tileTokensDim; - } else { - params.totalNumPaddedTokens = - params.ctasInTokenDimPerBatch * tileTokensDim * options.mNumBatches; - } - ctaOffset = maxNumCtas; - } else { - params.ptrTotalNumPaddedTokens = ptrTotalNumPaddedTokens; - params.ptrCtaIdxXyToBatchIdx = ptrCtaIdxXyToBatchIdx; - params.ptrCtaIdxXyToMnLimit = ptrCtaIdxXyToMnLimit; - ctaOffset = maxNumCtas; - } - - if (options.mUseDeepSeekFp8 && options.mDtypeC == tg::Dtype::E4m3) { - params.ptrDqSfsC = reinterpret_cast(dSfC); - } - - params.ptrA = ptrA; - params.ptrB = ptrB; - params.strideInBytesA = options.mK * tg::dtypeGetNumBits(options.mDtypeA) / 8; - params.strideInBytesB = options.mK * tg::dtypeGetNumBits(options.mDtypeB) / 8; - - params.ptrSfA = dSfA; - params.ptrSfB = dSfB; - params.ptrSfC = dSfC; - - // Do we pad A or B? - bool doPadA = tg::dtypeNeedsPadding(options.mDtypeA, options.mMmaKind, options.mMmaK, isSparseA); - bool doPadB = tg::dtypeNeedsPadding(options.mDtypeB, options.mMmaKind, options.mMmaK, isSparseA); - - if (!batchM) { - // A is the expert - if (0 != options.mM % options.mTileM) { - throw std::runtime_error("0 == mM %% tileM"); - } - params.tileStridePerBatch = options.mM / options.mTileM; - params.nm = options.mM; - // Shape/stride for gmem tensor A. - auto [shapeA, strideA, tileShapeA] = makeTmaShapeStrideAbc( - options, options.mM, options.mN, options.mK >> isSparseA, options.mTileM, options.mTileN, - options.mTileK >> isSparseA, MatrixType::MatrixA, options.mValidM, options.mValidN, - options.mValidK >> isSparseA); - // Build tma descriptor for A. - params.tmaA[0] = gemm::buildNdTmaDescriptor(options.mDtypeA, shapeA, strideA, tileShapeA, - const_cast(ptrA), doPadA, - /*doSwizzle=*/true); - - // The input is padded: - // [act0, padding, padding, ... TileN size .., act1, padding, padding, ...] - auto const inputNumTokens = ctaOffset * options.mTileN; - - if (!batchedGemm::doesRouteImplUseLdgsts(options.mRouteImpl)) { - bool useRouteAct = batchedGemm::doesRouteImplUseTma(options.mRouteImpl); - // B is the activation - // Shape/stride for gmem tensor B. - auto [shapeB, strideB, tileShapeB] = makeTmaShapeStrideAbc( - options, options.mM, useRouteAct ? options.mNumTokens : inputNumTokens, options.mK, - options.mTileM, (useRouteAct ? 1 : options.mTileN), options.mTileK, MatrixType::MatrixB, - options.mValidM, useRouteAct ? options.mNumTokens : inputNumTokens, options.mValidK); - // Build tma descriptor for B. - params.tmaB[0] = gemm::buildNdTmaDescriptor(options.mDtypeB, shapeB, strideB, tileShapeB, - const_cast(ptrB), doPadB, - /*doSwizzle=*/true); - } - - if (options.mDtypeA == tg::Dtype::E2m1 || options.mDtypeA == tg::Dtype::MxE4m3 || - options.mDtypeA == tg::Dtype::MxE2m1 || options.mDtypeA == tg::Dtype::MxInt4) { - tg::Dtype dTypeSfA{}; - if (options.mDtypeA == tg::Dtype::E2m1) { - dTypeSfA = tg::Dtype::E4m3; - } else if (options.mDtypeA == tg::Dtype::MxInt4) { - dTypeSfA = tg::Dtype::Bfloat16; - } else { - dTypeSfA = tg::Dtype::UE8m0; - } - - int32_t const numEltsPerSfA = options.mSfBlockSizeA; - - // Build TMA descriptor for gmem A block scaling factors. - auto [shapeSfA, strideSfA, tileShapesSfA] = makeTmaShapeStrideSfAb( - options.mM * options.mNumBatches, options.mN, options.mK, MatrixType::MatrixA, - options.mTileM, options.mTileN, options.mTileK, tg::SfLayout::R128c4, - options.mSfReshapeFactor, numEltsPerSfA); - params.tmaSfA[0] = gemm::buildSfTmaDescriptor(dTypeSfA, shapeSfA, strideSfA, tileShapesSfA, - const_cast(dSfA)); - } - - if (isSparseA) { - // Build TMA descriptor for gmem A sparsity information. - auto [shapeSparsityInfoA, strideSparsityInfoA, tileShapesSparsityInfoA] = - makeTmaShapeStrideSparsityInfoA(options); - params.tmaSparsityInfoA = - gemm::buildNdTmaDescriptor(tg::Dtype::UInt8, shapeSparsityInfoA, strideSparsityInfoA, - tileShapesSparsityInfoA, const_cast(ptrSparsityInfoA), - /*doPad=*/false, - /*doSwizzle=*/true); - } - - if (options.mDtypeB == tg::Dtype::E2m1 || options.mDtypeB == tg::Dtype::MxE4m3 || - options.mDtypeB == tg::Dtype::MxE2m1) { - tg::Dtype const dTypeSf = - (options.mDtypeB == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; - int32_t const numEltsPerSfB = options.mSfBlockSizeB; - - if (batchedGemm::doesRouteImplUseTma(options.mRouteSfsImpl.value())) { - // The input is NOT padded: - // [act0, act1, act2, ...] - - // Build TMA descriptor for gmem B block scaling factors. - // Pad number of scaling factors to the nearest multiple of 16 because of the TMA 16B - // alignment requirement. - auto numSfsInK = options.mK / numEltsPerSfB; - numSfsInK = ceilDiv(numSfsInK, 16) * 16; - - auto numSfsInValidK = options.mValidK / numEltsPerSfB; - numSfsInValidK = ceilDiv(numSfsInValidK, 16) * 16; - - auto [shapeSfB, strideSfB, tileShapesSfB] = makeTmaShapeStrideAbc( - options, options.mM, options.mNumTokens, numSfsInK, options.mTileM, 1 /* tileN */, - options.mTileK / numEltsPerSfB, MatrixType::MatrixB, options.mValidM, - options.mNumTokens, numSfsInValidK); - params.tmaSfB[0] = gemm::buildNdTmaDescriptor(dTypeSf, shapeSfB, strideSfB, tileShapesSfB, - const_cast(dSfB), - /*doPad=*/false, - /*doSwizzle=*/true); - } else if (batchedGemm::doesRouteImplUseNoRoute(options.mRouteSfsImpl.value())) { - // The input is padded: - // [act0, padding, padding, ... TileN size .., act1, padding, padding, ...] - - auto const inputNumTokensSfB = ctaOffset * options.mTileN; - - // Build TMA descriptor for gmem B block scaling factors. - auto [shapeSfB, strideSfB, tileShapesSfB] = - makeTmaShapeStrideSfAb(options.mM, inputNumTokensSfB, options.mK, MatrixType::MatrixB, - options.mTileM, options.mTileN, options.mTileK, - options.mSfLayoutB, options.mSfReshapeFactor, numEltsPerSfB); - params.tmaSfB[0] = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfB, strideSfB, tileShapesSfB, - const_cast(dSfB)); - } - } - - // C is the output activation - if (options.mUseTmaStore) { - // Shape/stride for gmem tensor C. - // NOTE: Output is *always* sanitized across the whole MNK range. This ensures maximum - // compatibility with the next BMM where unwritten part of the output could be polluted by - // NaNs. - auto [shapeC, strideC, tileShapeC] = makeTmaShapeStrideAbc( - options, options.mM, ctaOffset * options.mTileN, options.mK, options.mTileM, - options.mTileN, options.mTileK, MatrixType::MatrixC); - // Build tma descriptor for C. - params.tmaC[0] = - gemm::buildNdTmaDescriptor(options.mDtypeC, shapeC, strideC, tileShapeC, ptrC, - /*doPad=*/false); - - } else { - params.ptrC = ptrC; - } - - } else { - // B is the expert - if (0 != options.mN % options.mTileN) { - throw std::runtime_error("0 == mN %% tileN"); - } - params.tileStridePerBatch = options.mN / options.mTileN; - params.nm = options.mN; - // Shape/stride for gmem tensor B. - auto [shapeB, strideB, tileShapeB] = makeTmaShapeStrideAbc( - options, options.mM, options.mN, options.mK, options.mTileM, options.mTileN, options.mTileK, - MatrixType::MatrixB, options.mValidM, options.mValidN, options.mValidK); - // Build tma descriptor for B. - params.tmaB[0] = gemm::buildNdTmaDescriptor(options.mDtypeB, shapeB, strideB, tileShapeB, - const_cast(ptrB), doPadB, - /*doSwizzle=*/true); - - if (!batchedGemm::doesRouteImplUseLdgsts(options.mRouteImpl)) { - // A is the activation - // Shape/stride for gmem tensor A. - // The input is padded: - // [act0, padding, padding, ... tileM size .., act1, padding, padding, ...] - auto const inputNumTokens = ctaOffset * options.mTileM; - bool useRouteAct = batchedGemm::doesRouteImplUseTma(options.mRouteImpl); - auto [shapeA, strideA, tileShapeA] = makeTmaShapeStrideAbc( - options, useRouteAct ? options.mNumTokens : inputNumTokens, options.mN, options.mK, - useRouteAct ? 1 : options.mTileM, options.mTileN, options.mTileK, MatrixType::MatrixA, - useRouteAct ? options.mNumTokens : inputNumTokens, options.mValidN, options.mValidK); - // Build tma descriptor for A. - params.tmaA[0] = gemm::buildNdTmaDescriptor(options.mDtypeA, shapeA, strideA, tileShapeA, - const_cast(ptrA), doPadA, - /*doSwizzle=*/true); - } - - if (options.mDtypeA == tg::Dtype::E2m1 || options.mDtypeA == tg::Dtype::MxE4m3 || - options.mDtypeA == tg::Dtype::MxE2m1) { - tg::Dtype const dTypeSf = - (options.mDtypeA == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; - int32_t const numEltsPerSfA = options.mSfBlockSizeA; - - if (options.mRouteSfsImpl.value() == batchedGemm::RouteImpl::NoRoute) { - // The input is padded: - // [act0, padding, padding, ... tileM size .., act1, padding, padding, ...] - auto const inputNumTokensSfA = ctaOffset * options.mTileM; - - // Build TMA descriptor for gmem A block scaling factors. - auto [shapeSfA, strideSfA, tileShapesSfA] = - makeTmaShapeStrideSfAb(inputNumTokensSfA, options.mN, options.mK, MatrixType::MatrixA, - options.mTileM, options.mTileN, options.mTileK, - tg::SfLayout::R128c4, options.mSfReshapeFactor, numEltsPerSfA); - params.tmaSfA[0] = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfA, strideSfA, tileShapesSfA, - const_cast(dSfA)); - } - } - - if (options.mDtypeB == tg::Dtype::E2m1 || options.mDtypeB == tg::Dtype::MxE4m3 || - options.mDtypeB == tg::Dtype::MxE2m1) { - tg::Dtype const dTypeSf = - (options.mDtypeB == tg::Dtype::E2m1) ? tg::Dtype::E4m3 : tg::Dtype::UE8m0; - int32_t const numEltsPerSfB = options.mSfBlockSizeB; - - // Build TMA descriptor for gmem B block scaling factors. - auto [shapeSfB, strideSfB, tileShapesSfB] = makeTmaShapeStrideSfAb( - options.mM, options.mN * options.mNumBatches, options.mK, MatrixType::MatrixB, - options.mTileM, options.mTileN, options.mTileK, options.mSfLayoutB, - options.mSfReshapeFactor, numEltsPerSfB); - params.tmaSfB[0] = gemm::buildSfTmaDescriptor(dTypeSf, shapeSfB, strideSfB, tileShapesSfB, - const_cast(dSfB)); - } - - // C is the output activation - if (options.mUseTmaStore) { - // Shape/stride for gmem tensor C. - // NOTE: Output is *always* sanitized across the whole MNK range. This ensures maximum - // compatibility with the next BMM where unwritten part of the output could be polluted by - // NaNs. - auto [shapeC, strideC, tileShapeC] = makeTmaShapeStrideAbc( - options, ctaOffset * options.mTileM, options.mN, options.mK, options.mTileM, - options.mTileN, options.mTileK, MatrixType::MatrixC); - // Build tma descriptor for C. - params.tmaC[0] = - gemm::buildNdTmaDescriptor(options.mDtypeC, shapeC, strideC, tileShapeC, ptrC, - /*doPad=*/false); - } else { - params.ptrC = ptrC; - } - } - - params.k = options.mK; - params.numBatches = options.mNumBatches; - - params.rank = 0; - params.tpGrpSize = 1; - - params.ptrPartialRowMax = rowMax; - params.ptrRowMaxCompletionBars = rowMaxBars; - - params.ptrNumNonExitingCtas = ptrNumNonExitingCtas; - - // Set the per-token scale factors for MetaFP8 or scale inputs - params.ptrPerTokenSfA = ptrPerTokenSfA; - params.ptrPerTokenSfB = ptrPerTokenSfB; - params.ptrBias = ptrBias; - - return params; -} -#endif -}; // namespace KernelParamsSetup - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace batchedGemm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h deleted file mode 100644 index 24c98a58f9..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelParamsDecl.h +++ /dev/null @@ -1,545 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once -namespace batchedGemm { - -// This is device code - -struct KernelParams { - ////////////////////////////////////////////////////////////////////////////////////////////////// - // - // BatchedGemm parameters. - // - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Maximum number of CTAs in the batch-token dimension. - static constexpr int MaxNumCtas = 2048; - - // - // TMA descriptor for A. - // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from - // makeTmaShapeStrideAbc. - // - // If batchM: - // If batchStrideInTokens > 0: - // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), K]. - // Logical strides are [K, 1]. - // Tile box shape is [tileM, tileK]. - // Tile box strides are [tileK, 1]. - // Else // batchStrideInTokens == 0: - // Logical shape is [M, K]. - // Logical strides are [K, 1]. - // Tile box shape is [tileM, tileK]. - // Tile box strides are [tileK, 1]. - // - // If batchN: - // If layoutA is MatrixLayout::MajorK - // Logical shape is [B, divUpMul(M, tileM), K]. - // Logical strides are [divUpMul(M, tileM) * K, K, 1]. - // Tile box shape is [1, tileM, tileK]. - // Tile box strides are [0, tileK, 1]. - // If layoutA is MatrixLayout::Mn - // Logical shape is [B, K, divUpMul(M, tileM)]. - // Logical strides are [K * divUpMul(M, tileM), divUpMul(M, tileM), 1]. - // Tile box shape is [1, tileK, tileM]. - // Tile box strides are [0, tileM, 1]. - // If layoutA is MatrixLayout::BlockMajorK - // Logical shape is [B, K / blockK, divUpMul(M, tileM), blockK]. - // Logical strides are [K * divUpMul(M, tileM), divUpMul(M, tileM) * blockK, blockK, 1]. - // Tile box shape is [1, tileK / min(blockK, tileK), tileM, min(blockK, tileK)]. - // Tile box strides are [0, tileM * min(blockK, tileK), min(blockK, tileK), 1]. - // where blockK is 128B. - // - // Dtype is set from options.mDtypeA. - CUtensorMap tmaA[1]; - - // TMA descriptor for B. - // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from - // makeTmaShapeStrideAbc. - // - // If batchM: - // If layoutB is MatrixLayout::MajorK - // Logical shape is [B, divUpMul(N, tileN), K]. - // Logical strides are [divUpMul(N, tileN) * K, K, 1]. - // Tile box shape is [1, tileN, tileK]. - // Tile box strides are [0, tileK, 1]. - // If layoutB is MatrixLayout::MajorMn - // Logical shape is [B, K, divUpMul(N, tileN)]. - // Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN), 1]. - // Tile box shape is [1, tileK, tileN]. - // Tile box strides are [0, tileN, 1]. - // If layoutB is MatrixLayout::BlockMajorK - // Logical shape is [B, K / blockK, divUpMul(N, tileN), blockK]. - // Logical strides are [K * divUpMul(N, tileN), divUpMul(N, tileN) * blockK, blockK, 1]. - // Tile box shape is [1, tileK / min(blockK, tileK), tileN, min(blockK, tileK)]. - // Tile box strides are [0, tileN * min(blockK, tileK), min(blockK, tileK), 1]. - // where blockK is 128B. - // - // If batchN: - // If batchStrideInTokens > 0: - // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), K]. - // Logical strides are [K, 1]. - // Tile box shape is [tileN, tileK]. - // Tile box strides are [tileK, 1]. - // Else // batchStrideInTokens == 0: - // Logical shape is [N, K]. - // Logical strides are [K, 1]. - // Tile box shape is [tileN, tileK]. - // Tile box strides are [tileK, 1]. - // - // Dtype is set from options.mDtypeB. - CUtensorMap tmaB[1]; - - // TMA descriptor for C, (when useTmaStore is true) - // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from - // makeTmaShapeStrideAbc. - // - // If batchM: - // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B), N]. - // Logical strides are [N, 1]. - // Tile box shape is [epilogueTileM, epilogueTileN]. - // Tile box strides are [epilogueTileN, 1]. - // - // If batchN: - // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B), M]. - // Logical strides are [M, 1]. - // Tile box shape is [epilogueTileN, epilogueTileM]. - // Tile box strides are [epilogueTileM, 1]. - // - // Dtype is set from options.mDtypeC. - CUtensorMap tmaC[1]; - - // TMA descriptor for the block scaling factors for A, for MxFp{4,8} and NvFp4 formats. - // Must be setup using gemm::buildSfTmaDescriptor with shapes and strides from - // makeTmaShapeStrideSfAb. - // The layout of scaling factors for A is always R128c4. - // - // Let P be the number of elements per SF. P=16 for NvFp4, P=32 for Mx formats. - // M must be a multiple of 128. - // K must be a multiple of 4P. - // The "logical" shape is: [paddedM, K / P], where paddedM is - // sum(divUpMul(M[bi], tileM) for bi in B) if batchM, - // otherwise divUpMul(M, TileM) * B. - // The R128c4 layout is: [paddedM / 128, K / P / 4, 512]. - // The shape we use for TMA is: [paddedM / 128, K / P / 4, 2, 256]. - // - // Dtype is Dtype::E4m3 for NvFp4, Dtype::UE8m0 for Mx formats. - CUtensorMap tmaSfA[1]; - - // TMA descriptor for the block scaling factors for B, for MxFp{4,8} and NvFp4 formats. - // Must be setup using gemm::buildSfTmaDescriptor with shapes and strides from - // makeTmaShapeStrideSfAb. - // The layout of block scaling factors for B is controlled by options.mSfLayoutB. - // - // Let P be the number of elements per SF. P=16 for NvFp4, P=32 for Mx formats. - // The "logical" shape is: [paddedN, K / 16] - // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B) if batchN, - // otherwise divUpMul(N, TileN) * B. - // - // If the layout is R128c4, - // paddedN must be a multiple of 128. - // K must be a multiple of 4P. - // The R128c4 layout is: [paddedN / 128, K / P / 4, 512] - // The shape we use for TMA is: [paddedN / 128, K / P / 4, 2, 256] - // - // If the layout is R8c4, - // paddedN must be a multiple of 8. - // K must be a multiple of 4P. - // The R8c4 layout is: [paddedN / 8, K / P / 4, 32] - // The shape we use for TMA is: [paddedN / 8, K / P / 4 / repeats, repeats * 32] - // where repeats = min(tileK / P / 4, 8) - // - // Dtype is Dtype::E4m3 for NvFp4, Dtype::UE8m0 for Mx formats. - CUtensorMap tmaSfB[1]; - - // TMA descriptor for the sparsity information of A, if structured sparsity is used. - // Must be setup using gemm::buildNdTmaDescriptor with shapes and strides from - // makeTmaShapeStrideSparsityInfoA. - // - // When sparsityA is Any_2_4: - // 2 elements are non-zero in any chunk of 4 elements. - // A 4-bit index indicates the position of the non-zero elements. - // The shape in UInt8 is: [B, M, K / 8] - // - // When sparsityA is Pairwise_4_8: - // 4 elements are non-zero in any chunk of 8 elements. - // The zero and non-zero elements are grouped in pairs. - // A 4-bit index indicates the position of the non-zero pairs. - // The shape in UInt8 is: [B, M, K / 16] - // - // Dtype is Dtype::UInt8. - CUtensorMap tmaSparsityInfoA; - - // The input matrix A. - // If (routeAct == true && batchM), the shape is [M, K]. tmaA is not used. - // Otherwise, check layout of tmaA to see the shape and strides. - void const* ptrA{nullptr}; - - // The stride for matrix A in bytes. - // Equals to K * dtypeGetNumBits(dtypeA) / 8. - uint64_t strideInBytesA; - - // The input matrix B. - // If (routeAct == true && batchN), the shape is [N, K]. tmaB is not used. - // Otherwise, check layout of tmaB to see the shape and strides. - void const* ptrB{nullptr}; - // The stride for matrix B in bytes. - // Equals to K * dtypeGetNumBits(dtypeB) / 8. - uint64_t strideInBytesB; - - // The output matrix C. Check "logical" layout of tmaC to see the shape and strides. - void* ptrC{nullptr}; - - // Inputs and output are MxFp{4,8}, Fp8, NvFp4. - // The scaling factors to apply to the output - can be used to incorporate input scaling factors - // as described below: C = SEncC * act(SDecA * SDecB * A * Bl) . (SDecA * SDecB * A * Br) - // -> ScaleGate = SDecA * SDecB - // ScaleC = SDecA * SDecB * SEncC - // - // Only the inputs are MxFp{4,8}, Fp8, NvFp4. - // C = act(SDecA * SDecB * A * Bl) . (SDecA * SDecB * A * Br) - // -> ScaleGate = SDecA * SDecB - // ScaleC = SDecA * SDecB - // - // Only the output is MxFp{4,8}, Fp8, NvFp4. - // C = SEncC * act(A * Bl) . (A * Br) - // -> ScaleGate = 1 - // ScaleC = SEncC - // - // The output tensor scaling factor for MxFp{4,8}, Fp8, NvFp4 and DeepSeek FP8 quantization. - // TensorRT-LLM API requires a scaling factor on the device. - // Shape is [B]. One scaling factor per tensor in batch. - float const* ptrScaleC{nullptr}; - - // The pre-activation scaling factor (typically dequantA * dequantB) for non-gated non-linear - // activation. - // Only used when non-linear activation is applied (e.g., GELU, Relu2). - // When used, scaleC should be quantScaleC only, and this scale is applied before the - // activation. Shape is [B]. - float const* ptrScaleAct{nullptr}; - - // The output gate scale for MxFp{4,8}, Fp8, NvFp4 and DeepSeek FP8 quantization. - // TensorRT-LLM API requires a scaling factor on the device. - // Shape is [B]. One scaling factor per tensor in batch. - float const* ptrScaleGate{nullptr}; - - // The clamp limit before the activation. - // Shape is [B]. - // Clamp is INF if nullptr. - // If applied on SwiGlu, it will be: - // - // x_glu = x_glu.clamp(min=None, max=limit) - // x_linear = x_linear.clamp(min=-limit, max=limit) - float const* ptrClampLimit{nullptr}; - - // The alpha and beta for SwiGlu or GeGlu. - // Shape is [B]. One alpha and one beta per tensor in batch. - // Alpha is 1.f if nullptr. - // Beta is 0.f if nullptr. - // The formula for SwiGlu (for GeGlu, replace sigmoid with phi): - // - // out_glu = x_glu * torch.sigmoid(alpha * x_glu) * (x_linear + beta) - float const* ptrGatedActAlpha{nullptr}; - float const* ptrGatedActBeta{nullptr}; - - // The K dimension. It is the hidden dimension of the input matrices. - int32_t k; - - // The non-batched dimension. - // It is N if batchM, otherwise M. - int32_t nm; - - // Tile stride per batch for the non-batched dimension. - // It is N / TileN if batchM, otherwise M / TileM. - int32_t tileStridePerBatch; - - // TODO get rid of that. - // DeepSeek FP8 scaling factors for C - float* ptrDqSfsC{nullptr}; - - // The block scaling factors for A. - // The pointer must always be set regardless of the quantization recipe. - // If (routeAct == true && batchM), the shape is [M, K / 16]. tmaSfA is not used. - // For the layout (r128c4), see below. - // Otherwise, - // If MxFp{4,8} and NvFp4 formats are used, - // check the "logical" layout of tmaSfA to see the shape and strides. - // The dtype is Dtype::E4m3. - // - // If DeepSeek FP8 quantization recipe is used, - // If batchM: - // The shape is [K / 128, paddedM], - // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B). - // If batchN: - // The shape is [M / 128, K / 128], - // The rightmost dimension is contiguous in memory. - // The dtype is Dtype::Float32. - void const* ptrSfA{nullptr}; - - // The block scaling factors for B. - // The pointer must always be set regardless of the quantization recipe. - // If (routeAct == true && batchN), the shape is [N, K / 16]. tmaSfB is not used. - // For the layout (r128c4, r8c4), see below. - // Otherwise, - // If MxFp{4,8} and NvFp4 formats are used, - // check the layout of tmaSfB to see the shape and strides. - // The dtype is Dtype::E4m3. - // - // If DeepSeek FP8 quantization recipe is used, - // If batchM: - // The shape is [N / 128, K / 128], - // If batchN: - // The shape is [K / 128, paddedN], - // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). - // The rightmost dimension is contiguous in memory. - // The dtype is Dtype::Float32. - void const* ptrSfB{nullptr}; - - // The per-token scaling factors from scale A. - // - // This is used for either: - // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is Dtype::Float32 - // * When the routing scales are applied to the input activations (only when output is not - // transposed). The dtype is Dtype::Bfloat16 - // - // if (batchM (A is activations)): - // Logical shape is [sum(divUpMul(M[bi], tileM) for bi in B)] - // - // if (batchN (A is weights)): - // Logical shape is [B, divUpMul(M, tileM)] - // - void const* ptrPerTokenSfA{nullptr}; - - // The per-token scaling factors from scale B. - // - // This is used for either: - // * Per-token scaling factor quantization schemes, such as MetaFP8. The dtype is Dtype::Float32 - // * When the routing scales are applied to the input activations (only when output is - // transposed). The dtype is Dtype::Bfloat16 - // - // if (batchM (B is weights)): - // Logical shape is [B, divUpMul(N, tileN)] - // - // if (batchN (B is activations)): - // Logical shape is [sum(divUpMul(N[bi], tileN) for bi in B)] - void const* ptrPerTokenSfB{nullptr}; - - // The bias applied after the GEMM and before the activation function. - // The bias is applied before applying the global scaling factor. I.e. - // C = act(A * B + bias') * scaleC - // scaleC = dequantA * dequantB * quantC - // Thus, the bias' = bias / (dequantA * dequantB), where the bias is the original bias. - // - // If batchM, BiasType must be N, and bias shape is [B, N]. - // The bias is broadcasted along the M dimension. - // - // If batchNm BiasType must be M, and bias shape is [B, M]. - // The bias is broadcasted along the N dimension. - // - // The dtype is float32. - void const* ptrBias{nullptr}; - - // The output block scaling factors for C. - // - // If MxFp{4,8} and NvFp4 formats are used, - // The "logical" shape is: - // if batchM: [paddedM, N / 16] - // if batchN: [paddedN, M / 16] - // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B), - // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). - // - // If the layout is R128c4, - // paddedOuter must be a multiple of 128. - // inner must be a multiple of 64. - // The R128c4 layout is: [paddedOuter / 128, inner / 16 / 4, 512] - // The shape we use for TMA is: [paddedOuter / 128, inner / 16 / 4, 2, 256] - // where inner = N if batchM, otherwise M. - // where paddedOuter = paddedM if batchM, otherwise paddedN. - // - // If the layout is R8c4, - // paddedOuter must be a multiple of 8. - // inner must be a multiple of 64. - // The R8c4 layout is: [paddedOuter / 8, inner / 16 / 4, 32] - // The shape we use for TMA is: [paddedOuter / 8, inner / 16 / 4 / repeats, repeats * 32] - // where repeats = min(tileInner / 16 / 4, 8), - // where tileInner = tileN if batchM, otherwise tileM, - // where paddedOuter = paddedM if batchM, otherwise paddedN. - // where inner = N if batchM, otherwise M. - // - // The dtype is Dtype::E4m3. - // - // If DeepSeek FP8 quantization recipe is used, - // If batchM: - // The shape is [N / 128, paddedM], - // where paddedM is sum(divUpMul(M[bi], tileM) for bi in B). - // If batchN: - // The shape is [M / 128, paddedN], - // where paddedN is sum(divUpMul(N[bi], tileN) for bi in B). - // The rightmost dimension is contiguous in memory. - // The dtype is Dtype::Float32. - void* ptrSfC{nullptr}; - - ////////////////////////////////////////////////////////////////////////////////////////////////// - // - // Routing activations parameters. - // - ////////////////////////////////////////////////////////////////////////////////////////////////// - // These params are used when the kernel is configured with -routeAct true. - // The inputs are not padded, but the outputs are padded to divUpMul(M[bi], tileM) for batchM or - // divUpMul(N[bi], tileN) for batchN. - // If -routeAct is false, the params are not used and should be set to zero. - - // The routeMap for the input tokens. - // Map of expanded token index (counting the previous padded tokens) to the batch index - // the token belongs to. - // The shape is - // [sum(divUpMul(M[bi], tileM) for bi in B)] for batchM - // [sum(divUpMul(N[bi], tileN) for bi in B)] for batchN - // The dtype is int32_t. - // - // There are 3 tokens [0, 1, 2] such that [0, 1] belong to batch [B0] and [2] to batch [B1]. - // Let's assume that the padded size is 4. - // - // The expanded indices for tokens [0, 1, 2] are: - // expandedIdx[0] = 0 - // expandedIdx[1] = 1 - // expandedIdx[2] = divUpMul(2, 4) + 0 = 4 - // - // The route map is [B0, B0, X, X, B1, X, X, X] where X could be any value. - int32_t const* ptrRouteMap{nullptr}; - - // Total number of unpadded inputs - int32_t numTokens; - - // Total number of batches - int32_t numBatches; - - ////////////////////////////////////////////////////////////////////////////////////////////////// - // - // Batching information parameters. - // - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // In some cases, some CTAs must early-exit. E.g. when the grid size is set statically, but the - // actual workload is decided at runtime. This element on the device contains the number of CTAs - // that do not early-exit. The number corresponds to the X dim of the grid when the output is not - // transposed (i.e. batchM). To the Y dim, otherwise. - // The size is 1 and the dtype is int32_t. - // Used if isStaticBatch == false, otherwise set to nullptr. - // The pointer points to a scalar and the dtype is int32_t. The pointed value must be >= 0. - int32_t const* ptrNumNonExitingCtas{nullptr}; - - // Pointer to total number of padded tokens. - // Computed as - // int32_t totalNumPaddedTokens{0}; - // for (int bi = 0; bi < options.mNumBatches; bi++) { - // totalNumPaddedTokens += batchM ? divUpMul(options.mBatchedM[bi], options.mTileM) - // : divUpMul(options.mBatchedN[bi], options.mTileN); - // } - // The size is 1 and the dtype is int32_t. - // If isStaticBatch == true, ptrTotalNumPaddedTokens should be set to nullptr and - // totalNumPaddedTokens is used. - int32_t const* ptrTotalNumPaddedTokens{nullptr}; - - // Pointer to the map from the CTA index (in X/Y dim) to the batch index. - // Maps CTA index in batch dim (i.e. blockDim.x if batchM, otherwise blockDim.y) - // to batch index. - // E.g. with listM = 128,255,32 and tileM = 128, should be equal to - // ctaIdxXyToBatchIdx = [0, 1, 1, 2] - // If isStaticBatch == true, ptrCtaIdxXyToBatchIdx should be set to nullptr and ctaIdxXyToBatchIdx - // is used. - int32_t const* ptrCtaIdxXyToBatchIdx{nullptr}; - - // Pointer from the CTA index X/Y to the expanded tile index where the expanded tile index is - // computed as: - // - // int expandedIdx = 0; - // for (int bi = 0; bi < batchIdx-1; ++bi) { - // expandIdx = divUpMul(numTokens[bi], TileM/N); - // } - // expandIdx += - // E.g. with numTokens = [128,255,32] and tileM = 128, should be equal to - // ptrCtaIdxXyToMnLimit = [128, 256, 383, 416] - int32_t const* ptrCtaIdxXyToMnLimit{nullptr}; - - // Total number of padded tokens - used as the stride for the activation and C scaling factors. - // Check ptrTotalNumPaddedTokens to see how it is computed. - // If isStaticBatch == true, totalNumPaddedTokens is used, otherwise ptrTotalNumPaddedTokens. - int32_t totalNumPaddedTokens; - - // Total number of padded tokens - used as the stride for the output activation - // and C scaling factors. This is only used when isUniformNumTokensPerBatch is true. - int32_t totalNumOutputPaddedTokens; - - // A map from CTA index X/Y to batch index. - // Check ptrCtaIdxXyToBatchIdx to see how it is computed. - // If isStaticBatch == true, ctaIdxXyToBatchIdx is used, otherwise ptrCtaIdxXyToBatchIdx. - int32_t ctaIdxXyToBatchIdx[MaxNumCtas]; - - // **Expanded** limits for the batched dimension: - // tile * ctaIdxXyToTileIdxMn[ctaIdxXy] -> ctaIdxXyToMnLimit[ctaIdxXy] - // Check ptrCtaIdxXyToMnLimit to see how it is computed. - // If isStaticBatch == true, ctaIdxXyToMnLimit is used, otherwise ptrCtaIdxXyToMnLimit. - int32_t ctaIdxXyToMnLimit[MaxNumCtas]; - - // Total number of CTAs in the token dimension per batch. - // Used only when isUniformNumTokensPerBatch is true. - int32_t ctasInTokenDimPerBatch{0}; - - // Stride for the batched dimension in the number of CTAs. - // Used only when isUniformNumTokensPerBatch is true. - int32_t batchStrideInCtas{0}; - - ////////////////////////////////////////////////////////////////////////////////////////////////// - // - // All-reduce parameters. - // - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // The rank id of the current device in the multi-gpu space. - int rank; - // The number of peer devices in tensor-parallel group. - int tpGrpSize; - - ////////////////////////////////////////////////////////////////////////////////////////////////// - // - // GatedAct parameters. - // - ////////////////////////////////////////////////////////////////////////////////////////////////// - - // Pointer for partial row max for DeepSeek FP8 recipe. - // This is temporary storage for the row max results. - // If batchM, the shape is [2, totalNumPaddedTokens, N / 128] and the dtype is float. - // Otherwise, the shape is [2, totalNumPaddedTokens, M / 128] and the dtype is float. - float* ptrPartialRowMax{nullptr}; - - // Flags in global memory that sync on "exit" for row max computation. - // The shape is [numTilesM * numTilesN / 2] and the dtype is uint32_t, where - // if batchM, - // numTilesM = divUp(totalNumPaddedTokens, tileM). - // numTilesN = divUp(N, tileN). - // Otherwise, - // numTilesM = divUp(M, tileM). - // numTilesN = divUp(totalNumPaddedTokens, tileN). - // - // The memory must be set to 0 before the kernel launch. - uint32_t* ptrRowMaxCompletionBars{nullptr}; -}; - -/////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h deleted file mode 100644 index db7bf0e331..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/KernelTraits.h +++ /dev/null @@ -1,723 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include - -#include "Enums.h" -#include "trtllm/gen/CommonUtils.h" -#include "trtllm/gen/DtypeDecl.h" -#include "trtllm/gen/MmaDecl.h" -#include "trtllm/gen/SparsityDecl.h" - -namespace batchedGemm { - -namespace gemm { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace tg = trtllm::gen; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Structure to manage memory allocation with configurable reuse -class MemAllocatorHelper { - public: - // The default constructor. - MemAllocatorHelper() {} - - // Constructor to initialize chunk sizes, alignments, and reuse flags - MemAllocatorHelper(std::vector> const& sizes, - std::vector const& reuse, std::vector const& names) - : mNumBytesAndAlignmentPerSmemChunk(sizes), mFirstChunkReuse(reuse), mSmemChunkNames(names) {} - - // Function to calculate the size of the array from 0 to jj chunks - int32_t getOffsetBeforeChunk(int jj) const { - int32_t totalSize = 0; - for (int32_t ii = 0; ii < jj; ++ii) { - auto const& elem = mNumBytesAndAlignmentPerSmemChunk[ii]; - auto paddedSize = getSizePaddedToAlignment(elem.first, elem.second); - // If SMEM chunk is reused but the size of the current chunk is - // larger than currently counted size - if (mFirstChunkReuse[ii] && paddedSize > totalSize) { - // Set new size to the size of the current chunk. - // E.g. possible in case of - // mNumBytesAndAlignmentPerSmemChunk = {{1, 1}, {1, 1}, {1024, 1}} - // mFirstChunkReuse = {false, false, true} - // The last chunk is larger than the first plus second, so total size is 1024. - totalSize = paddedSize; - } else if (!mFirstChunkReuse[ii]) { - totalSize += paddedSize; - } - } - return totalSize; - } - - // Returns the offset of the ith chunk - int32_t getChunkOffsetByName(std::string const& name) const { - for (size_t ii = 0; ii < mSmemChunkNames.size(); ++ii) { - if (mSmemChunkNames[ii] == name) { - return getChunkOffset(ii); - } - } - throw std::runtime_error("Name not found: " + name); - } - - // Returns the first chunk reuse flag given chunk name. - int getFirstChunkReuseFlagByName(std::string const& name) const { - for (size_t ii = 0; ii < mSmemChunkNames.size(); ++ii) { - if (mSmemChunkNames[ii] == name) { - return getFirstChunkReuseFlag(ii); - } - } - throw std::runtime_error("Name not found: " + name); - } - - // Function to calculate the total size of the SMEM array - int32_t getTotalSize() const { - return getOffsetBeforeChunk(static_cast(mNumBytesAndAlignmentPerSmemChunk.size())); - } - - // Print the contents of this object. - void print() const { - for (size_t ii = 0; ii < mNumBytesAndAlignmentPerSmemChunk.size(); ++ii) { - printf("Chunk %zd %s: %d bytes, %d alignment, reuse %s, offset %d\n", ii, - mSmemChunkNames[ii].c_str(), mNumBytesAndAlignmentPerSmemChunk[ii].first, - mNumBytesAndAlignmentPerSmemChunk[ii].second, mFirstChunkReuse[ii] ? "true" : "false", - getChunkOffset(ii)); - } - } - - private: - int32_t getChunkOffset(int32_t ii) const { - if (mFirstChunkReuse[ii]) { - // Reuse the offset of the 0th chunk. - return getChunkOffset(0); - } - - // Get offset of ii chunks. - auto offset = getOffsetBeforeChunk(ii); - // Ensure alignment for the current chunk - return getSizePaddedToAlignment(offset, mNumBytesAndAlignmentPerSmemChunk[ii].second); - } - - // Returns the first chunk reuse flag for the ith chunk. - int getFirstChunkReuseFlag(int32_t ii) const { return mFirstChunkReuse[ii]; } - - // Helper function to calculate padded size - int32_t getSizePaddedToAlignment(int32_t size, int32_t alignment) const { - assert((alignment & (alignment - 1)) == 0); - return (size + alignment - 1) & ~(alignment - 1); - } - - private: - // Sizes and alignment requirements of each chunk - // NOTE: be careful and make sure that the memory dependency is clear and - // chunks in the beginning of the SMEM can be overwritten. - std::vector> mNumBytesAndAlignmentPerSmemChunk; - // Chunk reuse configuration. True at ith position means that ith chunk starts at smemOffset = 0. - std::vector mFirstChunkReuse; - // Buffer names for inspection purposes. - std::vector mSmemChunkNames; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int getNumSmemBitsPerElt(tg::Dtype dtype, tg::MmaKind mmaKind, int mmaK, bool isSparseA) { - if (mmaKind == tg::MmaKind::Auto) { - throw std::runtime_error("mmaKind != tg::MmaKind::Auto"); - } - if (mmaKind == tg::MmaKind::MxFp8Fp6Fp4) { - return 8; - } else { - return tg::dtypeGetNumBits(dtype); - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -class KernelTraits { - public: - // The default constructor. - KernelTraits() {} - - // The constructor. - KernelTraits(tg::Dtype dtypeA, tg::Dtype dtypeB, tg::Dtype dtypeC, tg::Dtype dtypeAcc, - tg::Dtype dtypeMmaA, tg::Dtype dtypeMmaB, tg::MmaKind mmaKind, - tg::Sparsity sparsityA, int32_t mmaK, int32_t tileM, int32_t tileN, int32_t tileK, - int32_t epilogueTileM, int32_t epilogueTileN, int32_t numEltsPerSfA, - int32_t numEltsPerSfB, int32_t numStages, int32_t numStagesMma, - int32_t numSlicesForSplitK, int32_t numSlicesForSliceK, SplitK splitK, - bool useTmaStore, bool transposeMmaOutput, AllReduceAlgo allReduceAlgo, - bool fuseUtccpWithUtcmma, bool useMaxTmemOverlap, int32_t numEpilogueWarps, - bool usePersistentScheduler, bool useDeepSeekFp8, bool usePerTokenSfA, - bool usePerTokenSfB, bool useTwoCtas, BiasType biasType) - : mMmaKind{mmaKind}, - mFuseUtccpWithUtcmma{fuseUtccpWithUtcmma}, - mUseMaxTmemOverlap{useMaxTmemOverlap}, - mNumEpilogueWarps{numEpilogueWarps} { - // - // SMEM - // - { - // [smemA ] (1024B aligned) - // [smemB ] (1024B aligned) - // [smemBShuffle ] (1024B aligned) - // [gmemC0 ] (1024B aligned) (if needed) - // [gmemC1 ] (1024B aligned) (if needed) - // [rowMax ] (16B aligned) (if needed) - // [sliceK ] (16B aligned) (if needed) - // [per-token SF ] (16B aligned) (if needed) - // [bias ] (16B aligned) (if needed) - // - // SMEM for smemA and smemB might be repurposed and used for gmemC0 and gmemC1: - // - // [..smemA..][..smemB..][..smemBShuffle..] - // [..gmemC0..][..gmemC1..][..rowMax..][..sliceK..][..per-token SF..][..bias..] - // - - if (mMmaKind == tg::MmaKind::Auto) { - mMmaKind = dtypeGetMmaKind(dtypeMmaA, dtypeMmaB); - } - - std::vector> numBytesAndAlignmentPerSmemChunk; - std::vector firstChunkReuseSmem; - // Buffer names for inspection purposes. - std::vector smemChunkNames; - - int const isSparseA = static_cast(tg::isSparse(sparsityA)); - - // LoadA - { - // Number of bytes in load A shared memory. - // If A is sparse, we load only the non-zero elements. - auto const numSmemBytesLoadA = numStages * tileM * (tileK >> isSparseA) * - getNumSmemBitsPerElt(dtypeA, mMmaKind, mmaK, isSparseA) / - 8 /* bits */; - // Number of bytes for load A alignment for TMA load. - auto const numBytesAlignmentLoadA = 1024; - // loadA is already at first chunk. No need to reuse it. - auto const reuseChunksSmemLoadA = false; - // Add info. - smemChunkNames.emplace_back("smemLoadA"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numSmemBytesLoadA, numBytesAlignmentLoadA)); - firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadA); - } - - // LoadB - { - // Number of bytes in load B shared memory. - auto const numSmemBytesLoadB = numStages * (useTwoCtas ? tileN / 2 : tileN) * tileK * - getNumSmemBitsPerElt(dtypeB, mMmaKind, mmaK, isSparseA) / - 8 /* bits */; - // Number of bytes for load B alignment for TMA load. - auto const numBytesAlignmentLoadB = 1024; - // No need to reuse the first chunk. - auto const reuseChunksSmemLoadB = false; - // Add info. - smemChunkNames.emplace_back("smemLoadB"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); - firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); - } - - // SmemBShuffle - // FIXME: we should be able either: - // - Do modification in-place. For that we need to resolve pipeline dependency between - // smemB -> shuffleSmemB -> mma - // - Do 4 TMA SW32 loads or several LDGSTS loads. - { - // Number of bytes in save shuffled B in shared memory. - auto const numSmemBytesLoadB = - numSlicesForSliceK > 1 - ? numStages * tileN * tileK * - getNumSmemBitsPerElt(dtypeB, mMmaKind, mmaK, isSparseA) / 8 /* bits */ - : 0; - // Number of bytes for load B alignment for TMA load. - auto const numBytesAlignmentLoadB = 1024; - // No need to reuse the first chunk. - auto const reuseChunksSmemLoadB = false; - - // Add info. - smemChunkNames.emplace_back("smemBShuffle"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numSmemBytesLoadB, numBytesAlignmentLoadB)); - firstChunkReuseSmem.emplace_back(reuseChunksSmemLoadB); - } - - // GmemC - // FIXME we might need to fix this for GemmGatedAct, it needs less SMEM to store gated output. - for (int resIdx = 0; resIdx < 2; ++resIdx) { - // Type of the data in the SMEM for GmemC - auto dtypeSmemC = dtypeC; - if (allReduceAlgo == AllReduceAlgo::TwoShot || numSlicesForSplitK > 1) { - dtypeSmemC = dtypeAcc; - } - // Smem is used for GmemC output tile for TMA store and SplitK in CGA. - bool usesSmemForGmemC = useTmaStore || doesSplitKUseDsmem(splitK); - // SMEM for at leader CTA in DSMEM split-k contains K slices. - auto extraGmemCMultiplier = doesSplitKUseDsmem(splitK) ? numSlicesForSplitK : 1; - if (numSlicesForSliceK > 1) { - // TileN is expanded in N dimension for slice-K. - extraGmemCMultiplier *= numSlicesForSliceK; - } - - if (resIdx != 0 && !useDeepSeekFp8) { - // No data for Epilogue1 in case of non-DeepSeek GEMM. - extraGmemCMultiplier = 0; - } - - if (numEpilogueWarps) { - extraGmemCMultiplier *= numEpilogueWarps / 4; - } - - // Number of bytes to store the output in smem. - auto const numBytesSmemStoreC = usesSmemForGmemC - ? extraGmemCMultiplier * epilogueTileM * epilogueTileN * - tg::dtypeGetNumBits(dtypeSmemC) / 8 /* bits */ - : 0; - // Number of bytes for store C alignment for TMA store. - auto const numBytesAlignmentStoreC = 1024; - // gmemC reuses loadAb memory for split-K in DSMEM. - // Epilogue1 does not reuse and continues after the memory allocated Epilogue0 - // NOTE: we can always reuse loadAb SMEM as long as we don't have persistent scheduler. - - auto const reuseFirstChunksSmemStoreC = - doesSplitKUseDsmem(splitK) && resIdx == 0 && !usePersistentScheduler; - - // Add info. - smemChunkNames.emplace_back("smemGmemC" + std::to_string(resIdx)); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numBytesSmemStoreC, numBytesAlignmentStoreC)); - firstChunkReuseSmem.emplace_back(reuseFirstChunksSmemStoreC); - } - - // SmemSparsityInfoA - { - // Number of bytes for sparsity info in SMEM. - auto const numBytesSmemSparsityInfoA = - numStages * tileM * tg::getNumBytesSparsityInfo(sparsityA, tileK); - // Number of bytes alignment for sparsity info in SMEM. - auto const numBytesAlignmentSparsityInfoA = 1024; - // No need to reuse the first chunk. - auto const reuseChunksSmemSparsityInfoA = false; - - // Add info. - smemChunkNames.emplace_back("smemSparsityInfoA"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numBytesSmemSparsityInfoA, numBytesAlignmentSparsityInfoA)); - firstChunkReuseSmem.emplace_back(reuseChunksSmemSparsityInfoA); - } - - // RowMax - { - // Number of dqSfsC per CTA. - auto const numDqSfsCPerCta = transposeMmaOutput ? tileM : tileN; - // Number of bytes for rowMax in SMEM. - auto const numBytesSmemRowMax = (useDeepSeekFp8 ? numDqSfsCPerCta : 0) * - tg::dtypeGetNumBits(tg::Dtype::Fp32) / 8 /* bits */; - // Number of bytes alignment for rowMax in SMEM. - auto const numBytesAlignmentRowMax = 16; - - // Add info. - smemChunkNames.emplace_back("smemRowMax"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numBytesSmemRowMax, numBytesAlignmentRowMax)); - firstChunkReuseSmem.emplace_back(false); - } - - // SliceK - { - // Real tile size before slice-K reduction. - auto const tileSize = - numSlicesForSliceK > 1 ? numSlicesForSliceK * tileM * numSlicesForSliceK * tileN : 0; - // Number of bytes for tile in SMEM. - auto const numBytesSmemTile = tileSize * tg::dtypeGetNumBits(dtypeAcc) / 8 /* bits */; - // Number of bytes alignment for rowMax in SMEM. - auto const numBytesAlignmentTile = 16; - - // Add info. - smemChunkNames.emplace_back("smemSliceK"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numBytesSmemTile, numBytesAlignmentTile)); - firstChunkReuseSmem.emplace_back(false); - } - - // Per-token Scale Factors - { - // Number of bytes for per-token scale factors - auto const numBytesSmemPerTokenSf = (usePerTokenSfA ? (tileM) * sizeof(float) : 0) + - (usePerTokenSfB ? (tileN) * sizeof(float) : 0); - // Number of bytes alignment for per-token scale factors - auto const numBytesAlignmentPerTokenSf = 16; - // Add info. - smemChunkNames.emplace_back("smemPerTokenSf"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numBytesSmemPerTokenSf, numBytesAlignmentPerTokenSf)); - firstChunkReuseSmem.emplace_back(false); - } - - // Bias - { - int32_t numBytesSmemBias = 0; - if (isBiasTypeN(biasType)) { - numBytesSmemBias = tileN * sizeof(float); - } else if (isBiasTypeM(biasType)) { - numBytesSmemBias = tileM * sizeof(float); - } else if (isBiasTypeMn(biasType)) { - numBytesSmemBias = tileM * tileN * sizeof(float); - } - // Number of bytes alignment for bias - auto const numBytesAlignmentBias = 16; - // Add info. - smemChunkNames.emplace_back("smemBias"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numBytesSmemBias, numBytesAlignmentBias)); - firstChunkReuseSmem.emplace_back(false); - } - - // Per-block absolute maximum for multi-warp reduction. - { - // Number of bytes: number of epilogue warps * number of tile columns. - auto const numBytesSmemBlockAmax = transposeMmaOutput ? 4 * tileN * sizeof(float) : 0; - // Number of bytes alignment. - auto const numBytesAlignmentBlockAmax = 16; - // Add info. - smemChunkNames.emplace_back("smemBlockAmax"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numBytesSmemBlockAmax, numBytesAlignmentBlockAmax)); - firstChunkReuseSmem.emplace_back(false); - } - - // SmemConstSfBuf - // A buffer used to copy constant values to TMEM. - { - // Do we need the buffer? - bool const useConstSfBuf = dtypeB == tg::Dtype::E4m3 && dtypeMmaB == tg::Dtype::MxE4m3; - // Number of bytes for the buffer. - auto const numSmemBytesConstSfBuf = useConstSfBuf ? 512 : 0; - // Number of bytes for the alignment of the buffer. - auto const numBytesAlignmentConstSfBuf = 16; - // No need to reuse the first chunk. - auto const reuseChunksSmemConstSfBuf = false; - - // Add info. - smemChunkNames.emplace_back("smemConstSfBuf"); - numBytesAndAlignmentPerSmemChunk.emplace_back( - std::make_pair(numSmemBytesConstSfBuf, numBytesAlignmentConstSfBuf)); - firstChunkReuseSmem.emplace_back(reuseChunksSmemConstSfBuf); - } - - // Create SMEM helper object. - mSmemAllocatorHelper = - MemAllocatorHelper(numBytesAndAlignmentPerSmemChunk, firstChunkReuseSmem, smemChunkNames); -#if 0 - // E.g., - // Chunk 0 smemLoadA: 32768 bytes, 1024 alignment, false, offset 0 - // Chunk 1 smemLoadB: 32768 bytes, 1024 alignment, false, offset 32768 - // Chunk 2 smemBShuffle: 0 bytes, 1024 alignment, false, offset 65536 - // Chunk 3 smemGmemC0: 65536 bytes, 1024 alignment, true, offset 0 - // Chunk 4 smemGmemC1: 65536 bytes, 1024 alignment, false, offset 65536 - // Chunk 5 smemRowMax: 512 bytes, 16 alignment, false, offset 131072 - // Chunk 6 smemSliceK: 0 bytes, 16 alignment, false, offset 131584 - // Chunk 7 smemPerTokenSf: 0 bytes, 16 alignment, false, offset 131584 - mSmemAllocatorHelper.print(); -#endif - } - - // - // TMEM - // - // [..D..][..A..][.SfA.][.SfB.] - { - std::vector> numBytesAndAlignmentPerTmemChunk; - std::vector firstChunkReuseTmem; - std::vector tmemChunkNames; - // Matrix D - { - // Two set of TMEM resources for D share epilogueTileN columns, - // | set0:epiTileN0 | set0:epiTileN1/set1:epiTileN0 | set1:epiTileN1 | - auto const numCols = mUseMaxTmemOverlap ? 2 * tileN - epilogueTileN : tileN; - // Number of columns for accumulators. - auto const numTmemColsD = numSlicesForSliceK * numCols * numStagesMma * - tg::dtypeGetNumBits(dtypeAcc) / - tg::dtypeGetNumBits(tg::Dtype::UInt32); - // Number of columns for D alignment. - auto const numColsAlignmentD = 2; - // No need to reuse TMEM. - auto const reuseChunksTmemD = false; - - // Add info. - tmemChunkNames.emplace_back("tmemD"); - numBytesAndAlignmentPerTmemChunk.emplace_back( - std::make_pair(numTmemColsD, numColsAlignmentD)); - firstChunkReuseTmem.emplace_back(reuseChunksTmemD); - } - - // Matrix A - { - // We use TMEM for A if we use slice-K or if we need to cast A. - bool const useTmemA = (numSlicesForSliceK > 1) || (dtypeMmaA != dtypeA); - // Number of columns for A. - auto const numTmemColsA = - useTmemA ? numStages * tileK / - (numSlicesForSliceK * tg::dtypeGetNumBits(tg::Dtype::UInt32) / - tg::dtypeGetNumBits(dtypeMmaA)) - : 0; - // Number of columns for A alignment. - auto const numColsAlignmentA = 4; - // No need to reuse TMEM. - auto const reuseChunksTmemA = false; - - // Add info. - tmemChunkNames.emplace_back("tmemA"); - numBytesAndAlignmentPerTmemChunk.emplace_back( - std::make_pair(numTmemColsA, numColsAlignmentA)); - firstChunkReuseTmem.emplace_back(reuseChunksTmemA); - } - - // Sf A - { - // Does the MMA require block scales in TMEM for A? - bool const useBlockScalingA = tg::dtypeIsBlockFmt(dtypeMmaA); - // Are the block scales constant? - bool const useConstSfA = useBlockScalingA && !tg::dtypeIsBlockFmt(dtypeA); - // TMEM cols group size in the K dimension. - int32_t kGroupSize = 4; - // Number of columns per stage. - int32_t const numColsPerStage = - useBlockScalingA ? ((tileK / (kGroupSize * numEltsPerSfA)) * - tg::getTmemColStridePerGroup(tileM, mmaK, kGroupSize)) - : 0; - // Number of columns for scaling factors of A. - auto const numTmemColsSfA = - useConstSfA ? tg::roundUp(numColsPerStage, 4) - : (numColsPerStage * (mFuseUtccpWithUtcmma ? 1 : numStages)); - // Number of columns for Sf alignment. - auto const numColsAlignmentSfA = 4; - // No need to reuse TMEM. - auto const reuseChunksTmemSfA = false; - - // Add info. - tmemChunkNames.emplace_back("tmemSfA"); - numBytesAndAlignmentPerTmemChunk.emplace_back( - std::make_pair(numTmemColsSfA, numColsAlignmentSfA)); - firstChunkReuseTmem.emplace_back(reuseChunksTmemSfA); - } - - // Sf B - { - // Does the MMA require block scales in TMEM for B? - bool const useBlockScalingB = tg::dtypeIsBlockFmt(dtypeMmaB); - // Are the block scales constant? - bool const useConstSfB = useBlockScalingB && !tg::dtypeIsBlockFmt(dtypeB); - // TMEM cols group size in the K dimension. - int32_t kGroupSize = 4; - // Number of columns per stage. - int32_t const numColsPerStage = - useBlockScalingB ? ((tileK / (kGroupSize * numEltsPerSfB)) * - tg::getTmemColStridePerGroup(tileN, mmaK, kGroupSize)) - : 0; - // Number of columns for scaling factors of B. - auto const numTmemColsSfB = - useConstSfB ? tg::roundUp(numColsPerStage, 4) - : (numColsPerStage * (mFuseUtccpWithUtcmma ? 1 : numStages)); - // Number of columns for Sf alignment. - auto const numColsAlignmentSfB = 4; - // No need to reuse TMEM. - auto const reuseChunksTmemSfB = false; - - // Add info. - tmemChunkNames.emplace_back("tmemSfB"); - numBytesAndAlignmentPerTmemChunk.emplace_back( - std::make_pair(numTmemColsSfB, numColsAlignmentSfB)); - firstChunkReuseTmem.emplace_back(reuseChunksTmemSfB); - } - - // Sparsity info for A - { - // Number of columns for the sparsity info for A (note: for Dense, this is 0). - auto const numTmemColsSparsityInfoA = - numStages * tg::getNumBytesSparsityInfo(sparsityA, tileK) / 4 /* bytes */; - // Number of columns for Sf alignment. - auto const numColsAlignmentSparsityInfoA = 2; - // No need to reuse TMEM. - auto const reuseChunksTmemSparsityInfoA = false; - - // Add info. - tmemChunkNames.emplace_back("tmemSparsityInfoA"); - numBytesAndAlignmentPerTmemChunk.emplace_back( - std::make_pair(numTmemColsSparsityInfoA, numColsAlignmentSparsityInfoA)); - firstChunkReuseTmem.emplace_back(reuseChunksTmemSparsityInfoA); - } - - // Create TMEM helper object. - mTmemAllocatorHelper = - MemAllocatorHelper(numBytesAndAlignmentPerTmemChunk, firstChunkReuseTmem, tmemChunkNames); - } - } - - public: - // The MMA kind. - tg::MmaKind mMmaKind; - // Whether fuse Utccp into the MMA task. - bool mFuseUtccpWithUtcmma; - // Whether use the max TMEM overlap trick. - bool mUseMaxTmemOverlap; - // The number of epilogue warps. - int32_t mNumEpilogueWarps; - // Helper for SMEM allocation. - MemAllocatorHelper mSmemAllocatorHelper; - // Helper for TMEM allocation. - MemAllocatorHelper mTmemAllocatorHelper; -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemBufferSize(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getTotalSize(); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getTmemBufferSize(KernelTraits traits) { - return traits.mTmemAllocatorHelper.getTotalSize(); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Starting address of each SMEM buffer. -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetLoadA(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemLoadA"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetLoadB(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemLoadB"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetLoadAb(KernelTraits traits) { return getSmemOffsetLoadA(traits); } - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetLoadShuffleB(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemBShuffle"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetGmemC(KernelTraits traits, int resIdx = 0) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemGmemC" + std::to_string(resIdx)); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetRowMax(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemRowMax"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetSliceK(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemSliceK"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetPerTokenSf(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemPerTokenSf"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetBias(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemBias"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetBlockAmax(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemBlockAmax"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetConstSfBuf(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemConstSfBuf"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getSmemOffsetSparsityInfoA(KernelTraits traits) { - return traits.mSmemAllocatorHelper.getChunkOffsetByName("smemSparsityInfoA"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t isSmemAbRepurposedToGmemC(KernelTraits traits, int resIdx = 0) { - return traits.mSmemAllocatorHelper.getFirstChunkReuseFlagByName("smemGmemC" + - std::to_string(resIdx)); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Starting address of each TMEM buffer. -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getTmemOffsetD(KernelTraits traits) { - return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemD"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getTmemOffsetA(KernelTraits traits) { - return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemA"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getTmemOffsetSfA(KernelTraits traits) { - return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemSfA"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getTmemOffsetSfB(KernelTraits traits) { - return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemSfB"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline int32_t getTmemOffsetSparsityInfoA(KernelTraits traits) { - return traits.mTmemAllocatorHelper.getChunkOffsetByName("tmemSparsityInfoA"); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gemm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h deleted file mode 100644 index b800ecaf83..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/TmaDescriptor.h +++ /dev/null @@ -1,295 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include - -#include "trtllm/gen/DtypeDecl.h" -#include "trtllm/gen/MmaDecl.h" - -#ifdef TLLM_ENABLE_CUDA -#include -#endif - -namespace batchedGemm { - -namespace gemm { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace tg = trtllm::gen; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -#ifdef TLLM_ENABLE_CUDA - -inline CUtensorMap buildNdTmaDescriptor(tg::Dtype dtype, std::vector const& shapes, - std::vector const& strides, - std::vector const& tileShapes, void* gmemAddr, - bool doPad, bool doSwizzle = true) { - // The multiplication factor of the data padding in SMEM. - int32_t padMultiplier = 1; - CUtensorMap desc{}; - // The data type. - CUtensorMapDataType tmaDataFormat{CU_TENSOR_MAP_DATA_TYPE_FLOAT32}; - if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::MxE4m3 || dtype == tg::Dtype::UE8m0 || - dtype == tg::Dtype::UInt8) { - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; - } else if (dtype == tg::Dtype::Fp16) { - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_FLOAT16; - } else if (dtype == tg::Dtype::Bfloat16) { - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; - } else if (dtype == tg::Dtype::E2m1) { - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; - } else if (dtype == tg::Dtype::MxE2m1 || dtype == tg::Dtype::MxInt4) { - if (doPad) { - padMultiplier = 2; - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B; - } else { - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B; - } - } else if (dtype == tg::Dtype::Fp32) { - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_FLOAT32; - } else { - std::cerr << "buildNdTmaDescriptor: unexpected dtype " << tg::dtypeToString(dtype) << std::endl; - assert(false); - } - - // The swizzle type. - CUtensorMapSwizzle swizzleType{CU_TENSOR_MAP_SWIZZLE_NONE}; - int32_t fastestDimTileSizeBytes = - (tileShapes[0] * tg::dtypeGetNumBits(dtype) * padMultiplier) / /* bits */ 8; - if (doSwizzle) { - if ((fastestDimTileSizeBytes % 128) == 0) { - swizzleType = CU_TENSOR_MAP_SWIZZLE_128B; - } else if ((fastestDimTileSizeBytes % 64) == 0) { - swizzleType = CU_TENSOR_MAP_SWIZZLE_64B; - } else if ((fastestDimTileSizeBytes % 32) == 0) { - swizzleType = CU_TENSOR_MAP_SWIZZLE_32B; - } else if ((fastestDimTileSizeBytes % 16) == 0 && - (dtype == tg::Dtype::UE8m0 || dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::E2m1 || - dtype == tg::Dtype::UInt8)) { - swizzleType = CU_TENSOR_MAP_SWIZZLE_NONE; - } else { - std::cerr << "buildNdTmaDescriptor: unexpected fastestDimTileSizeBytes " - << fastestDimTileSizeBytes << std::endl; - assert(false); - } - } - - // Check gmem address must be 16B-aligned - assert((reinterpret_cast(gmemAddr) & 0b1111) == 0); // - - // Check shape must be in range [1, 2^32] - int32_t dim = shapes.size(); - // Expect 2 dimensions for regular gemm, 3 dimensions for batched gemm or blocked layout, and 4 - // dimensions for batched gemm with blocked layout. - assert(dim == 2 || dim == 3 || dim == 4); - // Check shape range. - for (int32_t ii = 0; ii < dim; ++ii) { - assert(shapes[ii] >= (uint64_t(1))); // Size must be min 1 - assert(shapes[ii] <= (uint64_t(1) << 32)); // Size must be max 2^32 - } - - // TMA descriptor does not store the zeroth stride and assumes it is 1. - assert(static_cast(strides.size()) == dim); - assert(strides[0] == 1); - - // Build strides in bytes. - // cuTensorMapEncodeTiled ignores the stride of the first dimension (implicitly 1). - std::vector stridesInBytes(dim - 1); - for (int32_t ii = 0; ii < dim - 1; ++ii) { - stridesInBytes[ii] = (strides[ii + 1] * tg::dtypeGetNumBits(dtype)) / /* bits */ 8; - } - - // Set the number of elements in the packed uint32_t element. - auto const numEltsPerUInt32 = 4 * /* bits */ 8 / (tg::dtypeGetNumBits(dtype) * padMultiplier); - // The number of elements in 128B. - auto const numEltsIn128B = numEltsPerUInt32 /*4B*/ * 32; - // The number of tile K hidden size (per token) in each block of shared memory. - auto const numEltsInClampedFastestTileSize = std::min(numEltsIn128B, tileShapes[0]); - - // Build box dim array. If tileShapes is smaller than dim, just fill with 1s. - assert(static_cast(tileShapes.size()) <= dim); - std::vector boxDim(dim, 1); - boxDim[0] = numEltsInClampedFastestTileSize; - for (size_t ii = 1; ii < tileShapes.size(); ++ii) { - if (tileShapes[ii] > 256) { - std::cerr << "buildNdTmaDescriptor: boxDim too large " << tileShapes[ii] << std::endl; - assert(false); - } else { - boxDim[ii] = tileShapes[ii]; - } - } - - // Set tile strides to 1; - std::vector tileStrides(dim, 1); - - // Build the descriptor. - CUresult result = - cuTensorMapEncodeTiled(&desc, tmaDataFormat, - /*tensorRank=*/dim, gmemAddr, shapes.data(), stridesInBytes.data(), - boxDim.data(), tileStrides.data(), - /*interleave=*/CU_TENSOR_MAP_INTERLEAVE_NONE, swizzleType, - /*l2Promotion=*/CU_TENSOR_MAP_L2_PROMOTION_L2_128B, - /*oobFill=*/CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - - if (result != CUDA_SUCCESS) { - char const* errorString; - cuGetErrorString(result, &errorString); - std::stringstream ss; - ss << "Error: Failed to initialize the TMA descriptor. " << errorString << std::endl; - - ss << "tmaFormat: " << static_cast(tmaDataFormat) << " dim: " << dim - << " gmem: " << gmemAddr << std::endl; - - ss << "Shape: "; - for (int ii = 0; ii < dim; ++ii) { - ss << shapes[ii] << " "; - } - ss << std::endl; - - ss << "Stride: "; - for (int ii = 0; ii < dim - 1; ++ii) { - ss << stridesInBytes[ii] << " "; - } - ss << std::endl; - - ss << "tileShapes: "; - for (int ii = 0; ii < dim; ++ii) { - ss << boxDim[ii] << " "; - } - ss << std::endl; - - ss << "tileStrides: "; - for (int ii = 0; ii < dim; ++ii) { - ss << tileStrides[ii] << " "; - } - ss << std::endl; - ss << "swizzleType: " << int(swizzleType) << std::endl; - ss << "(in " << __FILE__ << ":" << __LINE__ << ")" << std::endl; - throw std::runtime_error(ss.str()); - } - - return desc; -} - -// TODO: make it work with the above descriptor? -inline CUtensorMap buildSfTmaDescriptor(tg::Dtype dtype, std::vector const& shapes, - std::vector const& strides, - const std::vector& tileShapes, void* gmemAddr) { - CUtensorMap desc{}; - CUtensorMapDataType tmaDataFormat{}; - if (dtype == tg::Dtype::E4m3 || dtype == tg::Dtype::UE8m0) { - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_UINT8; - } else if (dtype == tg::Dtype::Bfloat16) { - tmaDataFormat = CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; - } else { - std::cerr << "buildSfTmaDescriptor: unexpected dtype " << tg::dtypeToString(dtype) << std::endl; - assert(false); - } - - // No swizzle for scaling factors. - CUtensorMapSwizzle swizzleType = CU_TENSOR_MAP_SWIZZLE_NONE; - - // Check gmem address must be 16B-aligned - assert((reinterpret_cast(gmemAddr) & 0b1111) == 0); // - - // Check shape must be in range [1, 2^32] - int32_t dim = shapes.size(); - // Check shape range. - for (int32_t ii = 0; ii < dim; ++ii) { - assert(shapes[ii] >= (uint64_t(1))); // Size must be min 1 - assert(shapes[ii] <= (uint64_t(1) << 32)); // Size must be max 2^32 - } - - // TMA descriptor does not store the zeroth stride and assumes it is 1. - assert(static_cast(strides.size()) == dim); - assert(strides[0] == 1); - - // Build strides in bytes. - // cuTensorMapEncodeTiled ignores the stride of the first dimension (implicitly 1). - std::vector stridesInBytes(dim - 1); - for (int32_t ii = 0; ii < dim - 1; ++ii) { - stridesInBytes[ii] = (strides[ii + 1] * tg::dtypeGetNumBits(dtype)) / /* bits */ 8; - } - - // Set tile strides to 1; - std::vector tileStrides(dim, 1); - - // Build the descriptor. - CUresult result = cuTensorMapEncodeTiled(/*tensorMap=*/&desc, - /*tensorDataType=*/tmaDataFormat, - /*tensorRank=*/dim, - /*globalAddress=*/gmemAddr, - /*globalDim=*/shapes.data(), - /*globalStrides=*/stridesInBytes.data(), - /*boxDim=*/tileShapes.data(), - /*elementStrides=*/tileStrides.data(), - /*interleave=*/CU_TENSOR_MAP_INTERLEAVE_NONE, - /*swizzle=*/swizzleType, - /*l2Promotion=*/CU_TENSOR_MAP_L2_PROMOTION_L2_128B, - /*oobFill=*/CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); - - if (result != CUDA_SUCCESS) { - char const* errorString; - cuGetErrorString(result, &errorString); - std::stringstream ss; - ss << "Error: Failed to initialize the TMA descriptor for SF. " << errorString << std::endl; - - ss << "tmaFormat: " << static_cast(tmaDataFormat) << " dim: " << dim - << " gmem: " << gmemAddr << std::endl; - - ss << "shape:"; - for (uint32_t shape_i : shapes) { - ss << " " << shape_i; - } - ss << std::endl; - - ss << "stridesInBytes:"; - for (uint32_t stride_i : stridesInBytes) { - ss << " " << stride_i; - } - ss << std::endl; - - ss << "tileShapes:"; - for (uint32_t tileShape_i : tileShapes) { - ss << " " << tileShape_i; - } - ss << std::endl; - - ss << "tileStrides:"; - for (uint32_t tileStride_i : tileStrides) { - ss << " " << tileStride_i; - } - ss << std::endl; - - ss << "swizzleType: " << int(swizzleType) << std::endl; - ss << "(in " << __FILE__ << ":" << __LINE__ << ")" << std::endl; - throw std::runtime_error(ss.str()); - } - - return desc; -} - -#endif // defined TLLM_ENABLE_CUDA - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gemm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h deleted file mode 100644 index 53155c8ffb..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CommonUtils.h +++ /dev/null @@ -1,58 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -namespace batchedGemm { - -namespace trtllm { -namespace gen { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// -// TMA OOB optimization constants. -// -// CUDA Programming Guide states that "globalDim must be non-zero and less than or equal to 2^32". -// In practice, the kernel acts funny with TMA shape of 2^32 so we use 2^31. -constexpr unsigned long TmaDimMax = 1UL << 31; -// Chosen so that LargeN * XLargeN * sizeof(dtype) >= 2^64 which causes overflow and effectively -// becomes 0. As sizeof(dtype) can be as small as 0.5B, we choose LargeN = 2^30 and XLargeN = 2^35 -// so overflow can happen. -constexpr unsigned long LargeN = 1UL << 30; -// Used in TMA stride. Should be less than 2^40. -constexpr unsigned long XLargeN = 1UL << 35; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline T ceilDiv(T m, T n) { - return (m + n - T(1)) / n; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -template -inline T roundUp(T m, T n) { - return ceilDiv(m, n) * n; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gen -} // namespace trtllm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaArchDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaArchDecl.h deleted file mode 100644 index 2a1f2dcc78..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaArchDecl.h +++ /dev/null @@ -1,95 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Be careful when modifying this file as it is included by the generated kernels. For example, do -// not add TLLM_CHECK_* constructs in this file. Thanks! -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace batchedGemm { - -namespace trtllm { -namespace gen { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -enum class CudaArch { - // Hopper - Sm90a = 0, - // Blackwell - Sm100a, - // Blackwell-family - Sm100f, - // Blackwell Ultra - Sm103a, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool isArchHopper(CudaArch cudaArch) { return cudaArch == CudaArch::Sm90a; } - -inline bool isArchBlackwell(CudaArch cudaArch) { - return cudaArch == CudaArch::Sm100a || cudaArch == CudaArch::Sm100f || - cudaArch == CudaArch::Sm103a; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline std::string cudaArchToString(CudaArch cudaArch, bool isFull = true) { - switch (cudaArch) { - case CudaArch::Sm90a: - return isFull ? "90a" : "90"; - case CudaArch::Sm100a: - return isFull ? "100a" : "100"; - case CudaArch::Sm100f: - return isFull ? "100f" : "100"; - case CudaArch::Sm103a: - return isFull ? "103a" : "103"; - default: - assert(false); - return ""; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline CudaArch stringToCudaArch(std::string const& str) { - if (str == "90a") { - return CudaArch::Sm90a; - } else if (str == "100a") { - return CudaArch::Sm100a; - } else if (str == "100f") { - return CudaArch::Sm100f; - } else if (str == "103a") { - return CudaArch::Sm103a; - } else { - assert(false); - return CudaArch::Sm100a; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gen -} // namespace trtllm -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h deleted file mode 100644 index 42bc884f92..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/CudaKernelLauncher.h +++ /dev/null @@ -1,92 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#ifdef TLLM_ENABLE_CUDA -#include -#include - -#include -#include -#endif -namespace batchedGemm { - -namespace trtllm { -namespace gen { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -#ifdef TLLM_ENABLE_CUDA -inline CUresult launchKernel(void* kernelParams, void* cudaStream, int32_t smemSize, - CUfunction kernel, dim3 block3, dim3 grid3, dim3 cluster3, - bool enablesPdl) { - // Make sure we can launch with that much shared memory. - if (smemSize > 48 * 1024) { - CUresult result = - cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smemSize); - if (result != CUDA_SUCCESS) { - return result; - } - } - - auto clusterDim = cluster3.x * cluster3.y * cluster3.z; - - CUlaunchConfig launchConfig; - launchConfig.blockDimX = block3.x; - launchConfig.blockDimY = block3.y; - launchConfig.blockDimZ = block3.z; - launchConfig.gridDimX = grid3.x; - launchConfig.gridDimY = grid3.y; - launchConfig.gridDimZ = grid3.z; - launchConfig.hStream = reinterpret_cast(cudaStream); - launchConfig.sharedMemBytes = smemSize; - - CUlaunchAttribute launchAttrs[3]; - launchAttrs[0].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION; - launchAttrs[0].value.clusterDim.x = cluster3.x; - launchAttrs[0].value.clusterDim.y = cluster3.y; - launchAttrs[0].value.clusterDim.z = cluster3.z; - launchAttrs[1].id = CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE; - launchAttrs[1].value.clusterSchedulingPolicyPreference = - (clusterDim > 1) ? CU_CLUSTER_SCHEDULING_POLICY_SPREAD : CU_CLUSTER_SCHEDULING_POLICY_DEFAULT; - launchAttrs[2].id = CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION; - launchAttrs[2].value.programmaticStreamSerializationAllowed = enablesPdl; - launchConfig.attrs = launchAttrs; - launchConfig.numAttrs = 3; - - // Add setting for non-portable cluster size. - if (clusterDim > 8) { - CUresult result = - cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED, - 1 // Enable non-portable cluster sizes - ); - if (result != CUDA_SUCCESS) { - return result; - } - } - - // Launch the kernel. - return cuLaunchKernelEx(&launchConfig, kernel, &kernelParams, nullptr); -} -#endif - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gen -} // namespace trtllm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h deleted file mode 100644 index 76ea5ebc83..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/DtypeDecl.h +++ /dev/null @@ -1,287 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include -#include -#ifndef TLLM_GEN_EXPORT_INTERFACE -#include "trtllm/gen/MmaDecl.h" -#else -#include "MmaDecl.h" -#endif - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Be careful when modifying this file as it is included by the generated kernels. For example, do -// not add TLLM_CHECK_* constructs in this file. Thanks! -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace batchedGemm { - -namespace trtllm { -namespace gen { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -enum class Dtype : uint32_t { - -// We use the following encoding for the types: -// -// Byte 0: Identifier for the type (going from 0 to the number of data types - 1, -// Byte 1: Number of bits in the type, -// Byte 2: Bit 0: Is it an integer? 0x1 if true, 0x0 otherwise; -// Bit 4: is it signed? 0x1 if true, 0x0 otherwise. -// Byte 3: Is it a block format? 0x1 if true, 0x0 otherwise. - -#define TLLM_ENCODE_DTYPE(BlockFormatBit, SignedBit, IntegerBit, NumBits, Uid) \ - uint32_t { \ - (BlockFormatBit << 24) | (SignedBit << 20) | (IntegerBit << 16) | (NumBits << 8) | (Uid) \ - } - - // clang-format off - Bfloat16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 16u, /*uid*/ 0u), - Bool = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 1u, /*uid*/ 1u), - E2m1 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 4u, /*uid*/ 2u), - E2m3 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 6u, /*uid*/ 3u), - E3m2 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 6u, /*uid*/ 4u), - E4m3 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 5u), - E5m2 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 6u), - Fp16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 16u, /*uid*/ 7u), - Fp32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 32u, /*uid*/ 8u), - Int8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 9u), - Int32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 10u), - Int64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 11u), - MxE2m1 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 4u, /*uid*/ 12u), - MxE4m3 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 13u), - MxInt4 = TLLM_ENCODE_DTYPE(/*block*/ 1u, /*signed*/ 1u, /*int*/ 1u, /*bits*/ 4u, /*uid*/ 14u), - UE8m0 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 0u, /*bits*/ 8u, /*uid*/ 15u), - UInt8 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 8u, /*uid*/ 16u), - UInt16 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 16u, /*uid*/ 17u), - UInt32 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 32u, /*uid*/ 18u), - UInt64 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 64u, /*uid*/ 19u), - UInt128 = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 0u, /*int*/ 1u, /*bits*/ 128u, /*uid*/ 20u), - Void = TLLM_ENCODE_DTYPE(/*block*/ 0u, /*signed*/ 1u, /*int*/ 0u, /*bits*/ 0u, /*uid*/ 21u), -// clang-format on - -#undef TLLM_ENCODE_DTYPE -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// The number of bits in a data type? -inline int dtypeGetNumBits(Dtype dtype) { - constexpr uint32_t kMask = 0xffu << 8; - return static_cast((static_cast(dtype) & kMask) >> 8); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Does the format use block scaling? -inline bool dtypeIsBlockFmt(Dtype dtype) { - constexpr uint32_t kMask = 0xffu << 24; - return static_cast((static_cast(dtype) & kMask) >> 24); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Is a given data type a floating-point type? -inline bool dtypeIsFloat(Dtype dtype) { - constexpr uint32_t kMask = 0x1u << 16; - return dtype != Dtype::Void && 0 == (static_cast(dtype) & kMask); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Is a given data type an 8-bit floating-point type? -inline bool dtypeIsFp8(Dtype dtype) { return dtype == Dtype::E4m3 || dtype == Dtype::E5m2; } - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Is a given data type an integer type? -inline bool dtypeIsInt(Dtype dtype) { - constexpr uint32_t kMask = 0x1u << 16; - return (dtype != Dtype::Bool) && (0 != (static_cast(dtype) & kMask)); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Is a given data type signed? -inline bool dtypeIsSigned(Dtype dtype) { - constexpr uint32_t kMask = 0x1u << 20; - return (0 != (static_cast(dtype) & kMask)); -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// For logging and error reporting -inline std::string dtypeToString(Dtype dtype) { - switch (dtype) { - case Dtype::Bfloat16: - return "Bfloat16"; - case Dtype::Bool: - return "Bool"; - case Dtype::E2m1: - return "E2m1"; - case Dtype::E2m3: - return "E2m3"; - case Dtype::E3m2: - return "E3m2"; - case Dtype::E4m3: - return "E4m3"; - case Dtype::E5m2: - return "E5m2"; - case Dtype::Fp16: - return "Fp16"; - case Dtype::Fp32: - return "Fp32"; - case Dtype::Int8: - return "Int8"; - case Dtype::Int32: - return "Int32"; - case Dtype::Int64: - return "Int64"; - case Dtype::MxE4m3: - return "MxE4m3"; - case Dtype::MxE2m1: - return "MxE2m1"; - case Dtype::MxInt4: - return "MxInt4"; - case Dtype::UE8m0: - return "UE8m0"; - case Dtype::UInt8: - return "UInt8"; - case Dtype::UInt16: - return "UInt16"; - case Dtype::UInt32: - return "UInt32"; - case Dtype::UInt64: - return "UInt64"; - case Dtype::UInt128: - return "UInt128"; - case Dtype::Void: - return "Void"; - default: - assert(false); - return "Unsupported type"; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline Dtype dtypeEltType(Dtype dtype) { - switch (dtype) { - case Dtype::MxE2m1: - return Dtype::E2m1; - case Dtype::MxE4m3: - return Dtype::E4m3; - default: - return dtype; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Note: the block size from the options should be used instead. -// TODO: remove this function? -inline int dtypeNumEltsPerSf(Dtype dtype, bool useSparsity = false) { - switch (dtype) { - case Dtype::E2m1: - return useSparsity ? 32 : 16; - case Dtype::MxE2m1: - case Dtype::MxE4m3: - case Dtype::MxInt4: - return useSparsity ? 64 : 32; - default: - assert(false); - return -1; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Returns the dtype of scaling factors, if applicable. -inline Dtype dtypeGetBlockSfType(Dtype dtype) { - switch (dtype) { - case Dtype::E2m1: - return Dtype::E4m3; - case Dtype::MxE2m1: - case Dtype::MxE4m3: - return Dtype::UE8m0; - case Dtype::MxInt4: - return Dtype::Bfloat16; - default: - assert(false); - return Dtype::Void; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline MmaKind dtypeGetMmaKind(Dtype dtypeA, Dtype dtypeB) { - auto dtypeEltA = dtypeEltType(dtypeA); - auto dtypeEltB = dtypeEltType(dtypeB); - - // Note: the order of the conditions is important here. - if ((dtypeA == Dtype::Fp16 && dtypeB == Dtype::Fp16) || - (dtypeA == Dtype::Bfloat16 && dtypeB == Dtype::Bfloat16)) { - return MmaKind::Fp16; - } - - if ((dtypeA == Dtype::Int8 || dtypeA == Dtype::UInt8) && - (dtypeB == Dtype::Int8 || dtypeB == Dtype::UInt8)) { - return MmaKind::Int8; - } - - // This statement captures both MxE2m1 and E2m1. - if (dtypeEltA == Dtype::E2m1 && dtypeEltB == Dtype::E2m1) { - return MmaKind::MxFp4NvFp4; - } - - if ((dtypeA == Dtype::E4m3 || dtypeA == Dtype::E5m2 || dtypeA == Dtype::E2m3 || - dtypeA == Dtype::E3m2 || dtypeA == Dtype::E2m1) && - (dtypeB == Dtype::E4m3 || dtypeB == Dtype::E5m2 || dtypeB == Dtype::E2m3 || - dtypeB == Dtype::E3m2 || dtypeB == Dtype::E2m1)) { - return MmaKind::Fp8Fp6Fp4; - } - - // At this point we know that both dtypes are Mx types and not both MxE2m1 at the same time. - if ((dtypeEltA == Dtype::E4m3 || dtypeEltA == Dtype::E5m2 || dtypeEltA == Dtype::E2m3 || - dtypeEltA == Dtype::E3m2 || dtypeEltA == Dtype::E2m1) && - (dtypeEltB == Dtype::E4m3 || dtypeEltB == Dtype::E5m2 || dtypeEltB == Dtype::E2m3 || - dtypeEltB == Dtype::E3m2 || dtypeEltB == Dtype::E2m1)) { - return MmaKind::MxFp8Fp6Fp4; - } - return MmaKind::Tf32; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool dtypeNeedsPadding(Dtype dtype, MmaKind mmaKind, [[maybe_unused]] int mmaK, - [[maybe_unused]] bool isSparseA) { - bool needsPadding = mmaKind == MmaKind::MxFp8Fp6Fp4 && dtype == Dtype::MxE2m1; - return needsPadding; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gen -} // namespace trtllm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h deleted file mode 100644 index ba3275ec85..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/MmaDecl.h +++ /dev/null @@ -1,112 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include -#ifndef TLLM_GEN_EXPORT_INTERFACE -#include "trtllm/gen/CommonUtils.h" -#else // TLLM_GEN_EXPORT_INTERFACE -#include "CommonUtils.h" -#endif // TLLM_GEN_EXPORT_INTERFACE - -namespace batchedGemm { - -namespace trtllm { -namespace gen { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// The kind of the MMA instruction -enum class MmaKind : uint32_t { - // For Blackwell this follows the PTX ISA description of the MMA instructions. - // https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-kind-shapes - - // The MMA type is auto-detected from the dtypes of the input tensors - Auto = 0, - // Supports dtypeA = dtypeB = Fp16 and dtypeD = [Fp16, Fp32] - // or dtypeA = dtypeB = Bfloat16 and dtypeD = [Fp32] - // Corresponds to the kind::f16 of tcgen05.mma. - Fp16 = 1, - // Supports dtypeA/B = [E4m3, E5m2, E2m3, E3m2, E2m1] and dtypeD = [Fp16, Fp32] - // Corresponds to the kind::f8f6f4 of tcgen05.mma. - Fp8Fp6Fp4 = 2, - // Supports dtypeA = dtypeB = [Int8, Uint8] and dtypeD = [Int32] - // Corresponds to the kind::i8 of tcgen05.mma. - Int8 = 3, - // Supports dtypeA = dtypeB = [MxE2m1, E2m1] with block scale [UM8e0, UEm4e3] - // and dtypeD = [Fp32] - // Corresponds to the kind::mxf4nvf4 of tcgen05.mma. - MxFp4NvFp4 = 4, - // Supports dtype dtypeA = dtypeB = [MxE4m3, MxE2m1] with block scale [UM8e0] - // and dtypeD = [Fp32] - // Corresponds to the kind::mxf8f6f4 of tcgen05.mma. - MxFp8Fp6Fp4 = 5, - // Supports dtypeA = dtypeB = Tf32 with dtypeD = [Fp32] - // Corresponds to the kind::tf32 of tcgen05.mma. - Tf32 = 6 -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool mmaKindIsBlockFmt(MmaKind mmaKind) { - return mmaKind == MmaKind::MxFp8Fp6Fp4 || mmaKind == MmaKind::MxFp4NvFp4; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// For logging and error reporting -inline std::string mmaKindToString(MmaKind mmaKind) { - switch (mmaKind) { - case MmaKind::Auto: - return "Auto"; - case MmaKind::Fp16: - return "Fp16"; - case MmaKind::Fp8Fp6Fp4: - return "Fp8Fp6Fp4"; - case MmaKind::Int8: - return "Int8"; - case MmaKind::MxFp4NvFp4: - return "MxFp4NvFp4"; - case MmaKind::MxFp8Fp6Fp4: - return "MxFp8Fp6Fp4"; - case MmaKind::Tf32: - return "Tf32"; - default: - assert(false); - return "Unsupported type"; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Get the TMEM column stride per group (i.e. kGroupSize * blockSize K elements) -inline int32_t getTmemColStridePerGroup(int32_t tileMn, int32_t mmaK, int32_t kGroupSize) { - int32_t colStride = 2 * ceilDiv(tileMn, 64); - if (mmaK == 96) { - colStride = std::max(4, colStride); - } - return colStride; -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gen -} // namespace trtllm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h deleted file mode 100644 index 56b537ff42..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SfLayoutDecl.h +++ /dev/null @@ -1,94 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Be careful when modifying this file as it is included by the generated kernels. For example, do -// not add TLLM_CHECK_* constructs in this file. Thanks! -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace batchedGemm { - -namespace trtllm { -namespace gen { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// This enumeration defines layouts for storing scale factors for FP4, FP6, and FP8 formats. -enum class SfLayout { - // Scale factors are stored in the same order as the associated matrix. - // I.e., the SF buffer is a tensor [m, ⌈n/b⌉], where m, n, and b are respectively the number of - // rows, columns and the block size. - // The SF for the element (i, j) is stored at (i, j/b). - Linear = 0, - - // A tile of 8x4 is stored contiguously. The order of elements inside the tile, and the order - // of tiles, are both row-major. - // I.e., the SF buffer is a tensor [⌈m/8⌉, ⌈n/b/4⌉, 8, 4]. - // The SF for the element (i, j) is stored at (i/8, j/b/4, i%8, (j/b)%4). - R8c4, - - // A tile of 8x16 is stored contiguously. The order of elements inside the tile, and the order - // of tiles, are both row-major. - // I.e., the SF buffer is a tensor [⌈m/8⌉, ⌈n/b/16⌉, 8, 16]. - // The SF for the element (i, j) is stored at (i/8, j/b/16, i%8, (j/b)%16). - // - // NOTE: This is a niche format that is currently used for the weights of the - // LowLatency FP4 kernels. It is not meant as an interchange format. In - // addition to the above requirements it requires n to be a multiple of 256. - R8c16, - - // A tile of 128x4 is stored contiguously. Rows 0-31, 32-63, 64-95 and 96-127 are interleaved - // as illustrated below: - // | 0,0 | 0,1 | 0,2 | 0,3 | 32,0 | 32,1 | 32,2 | 32,3 | ... | 96,3 | - // | 1,0 | 1,1 | 1,2 | 1,3 | 33,0 | 33,1 | 33,2 | 33,3 | ... | 97,3 | - // | ... | ... | ... | ... | ... | ... | ... | ... | ... | ... | - // | 31,0 | 31,1 | 31,2 | 31,3 | 63,0 | 63,1 | 63,2 | 63,3 | ... | 127,3 | - // I.e., the SF buffer is a tensor [⌈m/128⌉, ⌈n/b/4⌉, 32, 4, 4] - // The SF for the element (i, j) is stored at (i/128, j/b/4, i%32, (i%128)/32, (j/b)%4). - R128c4, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline std::string sfLayoutToString(SfLayout layout) { - switch (layout) { - case SfLayout::Linear: - return "linear"; - case SfLayout::R8c4: - return "8x4"; - case SfLayout::R8c16: - return "8x16"; - case SfLayout::R128c4: - return "128x4"; - default: - assert(false); - return "Unsupported layout"; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gen -} // namespace trtllm - -} // namespace batchedGemm diff --git a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SparsityDecl.h b/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SparsityDecl.h deleted file mode 100644 index f093d5d1f9..0000000000 --- a/include/flashinfer/trtllm/batched_gemm/trtllmGen_bmm_export/trtllm/gen/SparsityDecl.h +++ /dev/null @@ -1,134 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & - * AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0 - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include - -//////////////////////////////////////////////////////////////////////////////////////////////////// -// -// Be careful when modifying this file as it is included by the generated kernels. For example, do -// not add TLLM_CHECK_* constructs in this file. Thanks! -// -//////////////////////////////////////////////////////////////////////////////////////////////////// - -namespace batchedGemm { - -namespace trtllm { -namespace gen { - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// This enumeration defines structured sparsity modes. Please refer to the PTX ISA for more details. -enum class Sparsity { - // No sparsity. - Dense, - - // For each chunk of 2 elements, 1 is non-zero. Only non-zero elements are stored. - // A 4-bit index is used to indicate the position of the non-zero element. - // The index may only take the value 0b1110 or 0b0100, other values are undefined behavior. - // - // 0b1110: 0b0100: - // |------ a ------|------ 0 ------| |------ 0 ------|------ a ------| - // | 11 | 10 | 01 | 00 | | 11 | 10 | 01 | 00 | - Any_1_2, - - // For each chunk of 4 elements, 2 are non-zero. Only non-zero elements are stored. - // A 4-bit index is used to indicate the position of the non-zero elements. - // Meaningful values are: 0b0100, 0b1000, 0b1100, 0b1001, 0b1101, 0b1110. - // Most other values are undefined behavior. - // - // E.g. 0b1100 corresponds to: - // |-- b --|-- 0 --|-- 0 --|-- a --| - // | 11 | 10 | 01 | 00 | - Any_2_4, - - // For each chunk of 8 elements, 4 are non-zero. Only non-zero elements are stored. - // Further, the zero and non-zero elements are grouped in pairs. - // A 4-bit index is used to indicate the position of the non-zero elements. - // Meaningful values are: 0b0100, 0b1000, 0b1100, 0b1001, 0b1101, 0b1110. - // Most other values are undefined behavior. - // - // E.g. 0b1100 corresponds to: - // | d | c | 0 | 0 | 0 | 0 | b | a | - // | 11 | 10 | 01 | 00 | - Pairwise_4_8, -}; - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline bool isSparse(Sparsity sparsity) { return sparsity != Sparsity::Dense; } - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -inline std::string sparsityToString(Sparsity sparsity) { - switch (sparsity) { - case Sparsity::Dense: - return "dense"; - case Sparsity::Any_1_2: - return "1:2"; - case Sparsity::Any_2_4: - return "2:4"; - case Sparsity::Pairwise_4_8: - return "4:8"; - default: - assert(false); - return "Unsupported sparsity"; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Size of a sparsity chunk, for sparse modes. -inline int32_t getSparsityChunkSize(Sparsity sparsity) { - switch (sparsity) { - case Sparsity::Any_1_2: - return 2; - case Sparsity::Any_2_4: - return 4; - case Sparsity::Pairwise_4_8: - return 8; - case Sparsity::Dense: - default: - assert(false); - return 0; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -// Number of bytes needed to store the sparsity information. -inline size_t getNumBytesSparsityInfo(Sparsity sparsity, size_t numElts) { - switch (sparsity) { - case Sparsity::Dense: - return 0; - case Sparsity::Any_1_2: - case Sparsity::Any_2_4: - case Sparsity::Pairwise_4_8: - return numElts / getSparsityChunkSize(sparsity) * 4 /*bits*/ / 8; - default: - assert(false); - return 0; - } -} - -//////////////////////////////////////////////////////////////////////////////////////////////////// - -} // namespace gen -} // namespace trtllm - -} // namespace batchedGemm