chore: import upstream snapshot with attribution

python/paddle/incubate/fp8/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.

python/paddle/incubate/fp8/deep_gemm/__init__.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepEP project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepEP/blob/main/LICENSE
+
+from . import jit  # noqa: F401
+from .jit_kernels import (  # noqa: F401
+    ceil_div,
+    gemm_fp8_fp8_bf16_nt,
+    get_col_major_tma_aligned_tensor,
+    get_m_alignment_for_contiguous_layout,
+    get_num_sms,
+    k_grouped_wgrad_gemm_fp8_fp8_fp32_nt,
+    m_grouped_gemm_fp8_fp8_bf16_nt_contiguous,
+    m_grouped_gemm_fp8_fp8_bf16_nt_masked,
+    set_num_sms,
+    wgrad_gemm_fp8_fp8_fp32_nt,
+)
+from .utils import calc_diff  # noqa: F401

python/paddle/incubate/fp8/deep_gemm/jit/__init__.py

@@ -0,0 +1,25 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from .compiler import (  # noqa: F401
+    NVCCCompiler,
+    NVRTCCompiler,
+    build,
+    get_nvcc_compiler,
+)
+from .runtime import FP8GemmRuntime, FP8WGradGemmRuntime, Runtime  # noqa: F401

python/paddle/incubate/fp8/deep_gemm/jit/compiler.py

@@ -0,0 +1,397 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from __future__ import annotations
+
+import functools
+import hashlib
+import os
+import re
+import subprocess
+import time
+import uuid
+from typing import Any
+
+import cuda.bindings
+from cuda.bindings import nvrtc
+
+from paddle.utils.cpp_extension.cpp_extension import CUDA_HOME
+
+from . import interleave_ffma
+from .runtime import Runtime, RuntimeCache
+
+runtime_cache = RuntimeCache()
+
+
+def hash_to_hex(s: str) -> str:
+    """Transforms input into a hexadecimal hash string"""
+    if isinstance(s, int):
+        return f"{s & 0xFFFFFFFFFFFFFFFF:016x}"
+    elif isinstance(s, str):
+        md5 = hashlib.md5()
+        md5.update(s.encode("utf-8"))
+        return md5.hexdigest()[0:12]
+    else:
+        raise TypeError(f"Unsupported type for hashing: {type(s)}")
+
+
+@functools.cache
+def get_jit_include_dir() -> str:
+    return f"{os.path.dirname(os.path.abspath(__file__))}"
+
+
+@functools.cache
+def get_deep_gemm_version() -> str:
+    md5 = hashlib.md5()
+
+    # Update include directories
+    include_dir = f"{get_jit_include_dir()}/../../../../include/paddle/fluid/fp8/deep_gemm/include"
+    assert os.path.exists(include_dir), (
+        f"Cannot find GEMM include directory {include_dir}"
+    )
+    for filename in filter(
+        lambda x: x.endswith(".cuh"), sorted(os.listdir(include_dir))
+    ):
+        with open(os.path.join(include_dir, filename), "rb") as f:
+            md5.update(f.read())
+
+    # Update `interleave_ffma.py`
+    with open(
+        os.path.join(
+            os.path.dirname(os.path.realpath(__file__)), "interleave_ffma.py"
+        ),
+        "rb",
+    ) as f:
+        md5.update(f.read())
+    return md5.hexdigest()[0:12]
+
+
+@functools.cache
+def get_nvcc_compiler() -> tuple[str, str]:
+    paths = []
+    if os.getenv("DG_JIT_NVCC_COMPILER"):
+        paths.append(os.getenv("DG_JIT_NVCC_COMPILER"))
+    paths.append(os.path.join(CUDA_HOME, "bin", "nvcc"))
+
+    # Try to find the first available NVCC compiler
+    least_version_required = "12.3"
+    version_pattern = re.compile(r"release (\d+\.\d+)")
+    for path in paths:
+        if os.path.exists(path):
+            command = [path, "--version"]
+            result = subprocess.run(
+                command,
+                capture_output=True,
+                text=True,
+            )
+            match = version_pattern.search(result.stdout)
+            version = match.group(1)
+            assert match, f"Cannot get the version of NVCC compiler {path}"
+            assert version >= least_version_required, (
+                f"NVCC {path} version {version} is lower than {least_version_required}"
+            )
+            return path, version
+    raise RuntimeError("Cannot find any available NVCC compiler")
+
+
+@functools.cache
+def get_default_user_dir():
+    if "DG_JIT_CACHE_DIR" in os.environ:
+        path = os.getenv("DG_JIT_CACHE_DIR")
+        os.makedirs(path, exist_ok=True)
+        return path
+    return os.path.join(os.path.expanduser("~"), ".deep_gemm")
+
+
+@functools.cache
+def get_tmp_dir():
+    return os.path.join(get_default_user_dir(), "tmp")
+
+
+@functools.cache
+def get_cache_dir():
+    return os.path.join(get_default_user_dir(), "cache")
+
+
+def make_tmp_dir():
+    tmp_dir = get_tmp_dir()
+    os.makedirs(tmp_dir, exist_ok=True)
+    return tmp_dir
+
+
+def put(path, data):
+    # Write and do POSIX atomic replace
+    tmp_file_path = os.path.join(
+        make_tmp_dir(), f"file.tmp.{uuid.uuid4()!s}.{hash_to_hex(path)}"
+    )
+    with open(tmp_file_path, "wb" if isinstance(data, bytes) else "w") as f:
+        f.write(data)
+    os.replace(tmp_file_path, path)
+
+
+class Compiler:
+    @classmethod
+    def signature(cls) -> str:
+        pass
+
+    @staticmethod
+    def __version__() -> tuple[int, int]:
+        pass
+
+    @classmethod
+    def compile(cls, name: str, code: str, target_path: str) -> None:
+        pass
+
+    @staticmethod
+    def flags() -> list[str]:
+        cpp_standard = int(os.getenv("DG_JIT_OVERRIDE_CPP_STANDARD", 20))
+        return [
+            f"-std=c++{cpp_standard}",
+            "--ptxas-options=--register-usage-level=10"
+            + (",--verbose" if "DG_JIT_PTXAS_VERBOSE" in os.environ else ""),
+            # Suppress some unnecessary warnings, such as unused variables for certain `constexpr` branch cases
+            "--diag-suppress=39,161,174,177,186,940",
+        ]
+
+    @staticmethod
+    def include_dirs() -> list[str]:
+        return [
+            f"{get_jit_include_dir()}/../../../../include/paddle/fluid/fp8/deep_gemm/include"
+        ]
+
+    @classmethod
+    def build(
+        cls,
+        name: str,
+        runtime_cls: type[Runtime],
+        kwargs: dict[str, Any] | None = None,
+    ) -> Runtime:
+        code = runtime_cls.generate(kwargs)
+        # Compiler flags
+        flags = cls.flags()
+
+        # Build signature
+        enable_sass_opt = cls.__version__() <= (12, 8) and not int(
+            os.getenv('DG_JIT_DISABLE_FFMA_INTERLEAVE', 0)
+        )
+        signature = f'{name}$${get_deep_gemm_version()}$${cls.signature()}$${flags}$${enable_sass_opt}$${code}'
+        name = f"kernel.{name}.{hash_to_hex(signature)}"
+        path = os.path.join(get_cache_dir(), name)
+
+        # Check runtime cache or file system hit
+        global runtime_cache
+        cached_runtime = runtime_cache.get(path, runtime_cls, name, kwargs)
+        if cached_runtime is not None:
+            if int(os.getenv("DG_JIT_DEBUG", 0)):
+                print(f"Using cached JIT runtime {name} during build")
+            return cached_runtime
+
+        # Compile into a temporary CU file
+        os.makedirs(path, exist_ok=True)
+        cubin_path = os.path.join(path, "kernel.cubin")
+        tmp_cubin_path = os.path.join(
+            make_tmp_dir(),
+            f"nvcc.tmp.{uuid.uuid4()!s}.{hash_to_hex(cubin_path)}.cubin",
+        )
+
+        start_time = time.time()
+        cls.compile(name, code, tmp_cubin_path)
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        if int(os.getenv("DG_JIT_DEBUG", 0)):
+            print(
+                f"Compilation of JIT runtime {name} took {elapsed_time:.2f} seconds."
+            )
+
+        # Interleave FFMA reuse
+        if enable_sass_opt:
+            interleave_ffma.process(tmp_cubin_path)
+
+        # Atomic replace files
+        os.replace(tmp_cubin_path, cubin_path)
+
+        # Put cache and return
+        runtime = runtime_cache.get(
+            path, runtime_cls, name, kwargs, force_enable_cache=True
+        )
+        assert runtime is not None
+        return runtime
+
+
+class NVCCCompiler(Compiler):
+    @staticmethod
+    def __version__() -> tuple[int, int]:
+        _, version = get_nvcc_compiler()
+        major, minor = map(int, version.split("."))
+        return major, minor
+
+    @classmethod
+    def signature(cls) -> str:
+        return f"{get_nvcc_compiler()[0]}+{cls.__version__()}"
+
+    @classmethod
+    def flags(cls) -> list[str]:
+        cxx_flags = [
+            "-fPIC",
+            "-O3",
+            "-fconcepts",
+            "-Wno-deprecated-declarations",
+            "-Wno-abi",
+        ]
+        return [
+            *super().flags(),
+            *[f"-I{d}" for d in cls.include_dirs()],
+            "-gencode=arch=compute_90a,code=sm_90a",
+            "-cubin",
+            "-O3",
+            "--expt-relaxed-constexpr",
+            "--expt-extended-lambda",
+            f"--compiler-options={','.join(cxx_flags)}",
+        ]
+
+    @classmethod
+    def compile(cls, name: str, code: str, target_path: str) -> None:
+        # Write the code
+        path = os.path.join(get_cache_dir(), name)
+        src_path = os.path.join(path, "kernel.cu")
+        put(src_path, code)
+        command = [
+            get_nvcc_compiler()[0],
+            src_path,
+            "-o",
+            target_path,
+            *cls.flags(),
+        ]
+        if int(os.getenv("DG_JIT_DEBUG", 0)) or int(
+            os.getenv("DG_JIT_PRINT_COMPILER_COMMAND", 0)
+        ):
+            print(f"Compiling JIT runtime {name} with command {command}")
+
+        result = subprocess.run(command, capture_output=True, text=True)
+        if result.returncode != 0:
+            print(
+                f"NVCC compilation failed: stdout: {result.stdout}, stderr: {result.stderr}"
+            )
+            raise AssertionError(f"Failed to compile {src_path}")
+
+
+class NVRTCCompiler(Compiler):
+    @staticmethod
+    def __version__() -> tuple[int, int]:
+        res, major, minor = nvrtc.nvrtcVersion()
+        if res != nvrtc.nvrtcResult.NVRTC_SUCCESS:
+            # Failed to get the actual NVRTC version, use cuda-bindings version instead
+            major, minor = map(int, cuda.bindings.__version__.split(".")[:2])
+        return major, minor
+
+    @classmethod
+    def signature(cls) -> str:
+        return f"nvrtc+{cls.__version__()}"
+
+    @staticmethod
+    def include_dirs() -> list[str]:
+        if CUDA_HOME is None:
+            raise RuntimeError("CUDA_HOME is required for NVRTC compilation")
+        return [get_jit_include_dir(), os.path.join(CUDA_HOME, "include")]
+
+    @classmethod
+    def flags(cls) -> list[str]:
+        flags = [
+            *super().flags(),
+            *[f"-I{d}" for d in cls.include_dirs()],
+            "--gpu-architecture=sm_90a",
+            "-default-device",
+        ]
+        # NOTES: PCH is vital for compilation speed
+        if cls.__version__() >= (12, 8):
+            flags += ["--pch"]
+            if int(os.getenv("DG_JIT_DEBUG", 0)):
+                flags += ["--pch-verbose=true"]
+        return flags
+
+    @classmethod
+    def compile(cls, name: str, code: str, target_path: str) -> None:
+        # Create program
+        code_bytes = bytes(code, "utf-8")
+        result, program = nvrtc.nvrtcCreateProgram(
+            code_bytes, bytes(name, "utf-8"), 0, [], []
+        )
+        assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, (
+            f"Failed to create program: {result}"
+        )
+
+        # Compile
+        options = [bytes(flag, "utf-8") for flag in cls.flags()]
+        if int(os.getenv("DG_JIT_DEBUG", 0)) or int(
+            os.getenv("DG_JIT_PRINT_COMPILER_COMMAND", 0)
+        ):
+            print(f"Compiling JIT runtime {name} with options: {options}")
+        compile_result = nvrtc.nvrtcCompileProgram(
+            program, len(options), options
+        )[0]
+
+        # Print compiler log
+        if (
+            int(os.getenv("DG_JIT_DEBUG", 0))
+            or compile_result != nvrtc.nvrtcResult.NVRTC_SUCCESS
+        ):
+            result, log_size = nvrtc.nvrtcGetProgramLogSize(program)
+            assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, (
+                f"Failed to get program log size: {result}"
+            )
+
+            log_bytes = bytes(log_size)
+            result = nvrtc.nvrtcGetProgramLog(program, log_bytes)[0]
+            assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, (
+                f"Failed to get program log: {result}"
+            )
+            print(f"Compiler log: {log_bytes.decode('utf-8')}")
+
+        # Exit if failed
+        assert compile_result == nvrtc.nvrtcResult.NVRTC_SUCCESS, (
+            f"Failed to compile program: {compile_result}"
+        )
+
+        # Create CUBIN
+        result, cubin_size = nvrtc.nvrtcGetCUBINSize(program)
+        assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, (
+            f"Failed to get CUBIN size: {result}"
+        )
+        cubin_bytes = bytes(cubin_size)
+        result = nvrtc.nvrtcGetCUBIN(program, cubin_bytes)[0]
+        assert result == nvrtc.nvrtcResult.NVRTC_SUCCESS, (
+            f"Failed to get CUBIN: {result}"
+        )
+
+        # Write into the file system
+        put(target_path, cubin_bytes)
+
+        # Destroy handler
+        assert (
+            nvrtc.nvrtcDestroyProgram(program)[0]
+            == nvrtc.nvrtcResult.NVRTC_SUCCESS
+        ), f"Failed to destroy program: {result}"
+
+
+def build(
+    name: str, runtime_cls: type[Runtime], kwargs: dict[str, Any] | None = None
+) -> Runtime:
+    compiler_cls = (
+        NVRTCCompiler if int(os.getenv("DG_JIT_USE_NVRTC", 0)) else NVCCCompiler
+    )
+    return compiler_cls.build(name, runtime_cls, kwargs)

python/paddle/incubate/fp8/deep_gemm/jit/interleave_ffma.py

@@ -0,0 +1,167 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+import argparse
+import mmap
+import os
+import re
+import subprocess
+
+from paddle.utils.cpp_extension.cpp_extension import CUDA_HOME
+
+
+def run_cuobjdump(file_path):
+    command = [f"{CUDA_HOME}/bin/cuobjdump", "-sass", file_path]
+    result = subprocess.run(command, capture_output=True, text=True)
+    assert result.returncode == 0
+    return result.stdout
+
+
+def extract_ffma(sass):
+    lines = sass.splitlines()
+    collected = []
+    current = []
+
+    arch_name, func_name = "N/A", "N/A"
+    skip_next_line = False
+    for line in lines:
+        if "code for" in line:
+            arch_name = line.lstrip().replace("code for ", "", 1).rstrip()
+        elif "Function :" in line:
+            func_name = line.lstrip().replace("Function :", "", 1).rstrip()
+        elif "FFMA" in line:
+            current.append(line)
+            skip_next_line = True
+        elif skip_next_line:
+            current.append(line)
+            skip_next_line = False
+        else:
+            if len(current) >= 16:
+                assert len(current) % 2 == 0
+                collected.append((f"{arch_name}::{func_name}", current))
+            current = []
+
+    if int(os.getenv("DG_JIT_PRINT_REG_REUSE", 0)):
+        print(f"Found {len(collected)} FFMA segments")
+    return collected
+
+
+def extract_hex_from_line(line):
+    match = re.search(r"/\*\s*(0x[0-9a-fA-F]+)\s*\*/", line)
+    assert match
+    return int(match.group(1), 16)
+
+
+def validate(m, offset, le_bytes, num_lines):
+    assert len(le_bytes) == num_lines // 2
+    assert m[offset : offset + 16] == le_bytes[0]
+    for i in range(1, num_lines // 2):
+        if m[offset + i * 16 : offset + i * 16 + 16] != le_bytes[i]:
+            return False
+    return True
+
+
+def parse_registers(line):
+    line = re.sub(r"/\*.*?\*/", "", line)
+    line = line.replace(";", "")
+    tokens = line.strip().split(",")
+    registers = []
+    for token in tokens:
+        token = token.strip()
+        words = token.split()
+        for word in words:
+            if word.startswith("R"):
+                reg = word.split(".")[0]
+                registers.append(reg)
+    return registers
+
+
+def modify_segment(m, name, ffma_lines):
+    num_lines = (len(ffma_lines) * 9 // 16) // 2 * 2
+    assert num_lines % 2 == 0
+
+    le_bytes, new_le_bytes = [], []
+    reused_list = []
+    dst_reg_set = set()
+    last_reused, last_dst_reg = False, ""
+    num_changed = 0
+    for i in range(num_lines // 2):
+        dst_reg = parse_registers(ffma_lines[i * 2])[-2]
+        low_line, high_line = ffma_lines[i * 2], ffma_lines[i * 2 + 1]
+        low_hex, high_hex = (
+            extract_hex_from_line(low_line),
+            extract_hex_from_line(high_line),
+        )
+        le_bytes.append(
+            low_hex.to_bytes(8, "little") + high_hex.to_bytes(8, "little")
+        )
+        reused = (high_hex & 0x0800000000000000) != 0
+        if reused:
+            is_first_occurred = dst_reg not in dst_reg_set
+            if is_first_occurred or (last_reused and dst_reg == last_dst_reg):
+                # Modify the `reuse` and `yield` bits
+                assert high_hex & 0x0800200000000000, f"{hex(high_hex)}"
+                high_hex ^= 0x0800200000000000
+                reused = False
+                num_changed += 1
+            else:
+                reused_list.append(i)
+        dst_reg_set.add(dst_reg)
+        new_le_bytes.append(
+            low_hex.to_bytes(8, "little") + high_hex.to_bytes(8, "little")
+        )
+        last_reused, last_dst_reg = reused, dst_reg
+    if int(os.getenv("DG_JIT_PRINT_REG_REUSE", 0)):
+        print(
+            f" > segment `{name}` new reused list ({num_changed} changed): {reused_list}"
+        )
+
+    # Find the offset
+    offsets = []
+    offset = m.find(le_bytes[0])
+    while offset != -1:
+        offsets.append(offset)
+        offset = m.find(le_bytes[0], offset + 1)
+    offsets = list(
+        filter(lambda x: validate(m, x, le_bytes, num_lines), offsets)
+    )
+
+    # Replace with `new_le_bytes`
+    for offset in offsets:
+        for i in range(num_lines // 2):
+            m[offset + i * 16 : offset + i * 16 + 16] = new_le_bytes[i]
+
+
+def process(path):
+    if int(os.getenv("DG_JIT_PRINT_REG_REUSE", 0)):
+        print(f"Processing {path}")
+    output = run_cuobjdump(path)
+    segments = extract_ffma(output)
+    with open(path, "r+b") as f:
+        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_WRITE)
+        for segment in segments:
+            modify_segment(mm, *segment)
+        mm.close()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Interleave FFMA reg reuse")
+    parser.add_argument("--so", help="Path to the SO file")
+    args = parser.parse_args()
+
+    process(args.so)

python/paddle/incubate/fp8/deep_gemm/jit/runtime.py

@@ -0,0 +1,447 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+from __future__ import annotations
+
+import ctypes
+import enum
+import os
+import subprocess
+import time
+from typing import Any
+
+import cuda.bindings.driver as cbd
+
+import paddle
+from paddle.utils.cpp_extension.cpp_extension import CUDA_HOME
+
+
+def get_num_math_warpgroups(block_m: int) -> int:
+    return 1 if block_m == 64 else 2
+
+
+def get_num_threads_per_sm(
+    num_tma_threads: int, num_math_threads_per_group: int, block_m: int
+) -> int:
+    assert num_math_threads_per_group == 128, (
+        "Only support 128 threads per math group"
+    )
+    return (
+        get_num_math_warpgroups(block_m) * num_math_threads_per_group
+        + num_tma_threads
+    )
+
+
+class GemmType(enum.Enum):
+    Normal = 0
+    GroupedContiguous = 1
+    GroupedMasked = 2
+
+    def __str__(self) -> str:
+        return {
+            0: "Normal",
+            1: "GroupedContiguous",
+            2: "GroupedMasked",
+        }[self.value]
+
+
+class Runtime:
+    def __init__(self, path: str) -> None:
+        self.path = path
+        self.lib = None
+        self.kernel = None
+        assert self.is_path_valid(self.path)
+
+    @staticmethod
+    def is_path_valid(path: str) -> bool:
+        # Exists and is a directory
+        if not os.path.exists(path) or not os.path.isdir(path):
+            return False
+
+        # Contains all necessary files
+        files = ["kernel.cubin"]
+        return all(os.path.exists(os.path.join(path, file)) for file in files)
+
+    @staticmethod
+    def generate(kwargs: dict[str, Any]) -> str:
+        raise NotImplementedError
+
+    @staticmethod
+    def launch(kernel: cbd.CUkernel, kwargs: dict[str, Any]) -> cbd.CUresult:
+        raise NotImplementedError
+
+    def __call__(self, **kwargs) -> cbd.CUresult:
+        # Load CUBIN
+        if self.kernel is None:
+            start_time = time.time_ns()
+
+            # Load CUBIN
+            path = bytes(os.path.join(self.path, "kernel.cubin"), "utf-8")
+            result, self.lib = cbd.cuLibraryLoadFromFile(
+                path, [], [], 0, [], [], 0
+            )
+            assert result == cbd.CUresult.CUDA_SUCCESS, (
+                f"Failed to load library: {result}"
+            )
+
+            # Extract the kernel name
+            # TODO: use `cuda-bindings` API to do this (requires at least 12.8)
+            command = [f"{CUDA_HOME}/bin/cuobjdump", "-symbols", path]
+            result = subprocess.run(
+                command,
+                capture_output=True,
+                text=True,
+            )
+            assert result.returncode == 0
+            illegal_names = [
+                "vprintf",
+                "__instantiate_kernel",
+                "__internal",
+                "__assertfail",
+            ]
+            check_illegal = lambda line: any(
+                name in line for name in illegal_names
+            )
+            kernel_names = [
+                line.split()[-1]
+                for line in result.stdout.splitlines()
+                if line.startswith("STT_FUNC") and not check_illegal(line)
+            ]
+            assert len(kernel_names) == 1, (
+                f"Too many kernels in the library: {kernel_names}"
+            )
+
+            # Load kernel from the library
+            result, self.kernel = cbd.cuLibraryGetKernel(
+                self.lib, bytes(kernel_names[0], encoding="utf-8")
+            )
+            assert result == cbd.CUresult.CUDA_SUCCESS, (
+                f"Failed to load kernel: {result}"
+            )
+
+            end_time = time.time_ns()
+            elapsed_time = (end_time - start_time) / 1e6
+            if int(os.getenv("DG_JIT_DEBUG", 0)):
+                print(
+                    f"Loading JIT runtime {self.path} took {elapsed_time:.2f} ms."
+                )
+
+        # noinspection PyArgumentList
+        return self.launch(self.kernel, kwargs)
+
+    def __del__(self) -> None:
+        if self.lib is not None:
+            res = cbd.cuLibraryUnload(self.lib)[0]
+            if res != cbd.CUresult.CUDA_SUCCESS:
+                raise Exception(f"Failed to unload library {self.path}: {res}")
+
+
+class RuntimeCache:
+    def __init__(self) -> None:
+        self.cache = {}
+
+    def __setitem__(self, path: str, runtime: Runtime) -> None:
+        self.cache[path] = runtime
+
+    def get(
+        self,
+        path: str,
+        runtime_cls: type[Runtime],
+        name: str = "",
+        kwargs: dict[str, Any] | None = None,
+        force_enable_cache: bool = False,
+    ) -> Runtime | None:
+        # In Python runtime
+        if path in self.cache:
+            return self.cache[path]
+
+        # Already compiled
+        use_cache = force_enable_cache or not int(
+            os.getenv("DG_JIT_DISABLE_CACHE", 0)
+        )
+        if use_cache and os.path.exists(path) and Runtime.is_path_valid(path):
+            # Print heuristic for the first time
+            if name and (
+                int(os.getenv("DG_JIT_DEBUG", 0))
+                or int(os.getenv("DG_PRINT_CONFIGS", 0))
+            ):
+                simplified_kwargs = {}
+                for key, value in (
+                    kwargs.items() if kwargs is not None else {}.items()
+                ):
+                    value = (
+                        f"paddle.Tensor<{value.dtype}>"
+                        if isinstance(value, paddle.Tensor)
+                        else value
+                    )
+                    value = (
+                        "cuda.bindings.driver.CUtensorMap"
+                        if isinstance(value, cbd.CUtensorMap)
+                        else value
+                    )
+                    simplified_kwargs[key] = value
+                print(
+                    f"Put kernel {name} with {simplified_kwargs} into runtime cache"
+                )
+
+            runtime = runtime_cls(path)
+            self.cache[path] = runtime
+            return runtime
+        return None
+
+
+def get_cache_key(kwargs, num_tma_threads, num_math_threads_per_group):
+    key_params = {
+        "NUM_TMA_MULTICAST": kwargs["NUM_TMA_MULTICAST"],
+        "NUM_SMS": kwargs["NUM_SMS"],
+        "BLOCK_M": kwargs["BLOCK_M"],
+        "SMEM_SIZE": kwargs["SMEM_SIZE"],
+        "STREAM": kwargs["STREAM"],
+        "num_tma_threads": num_tma_threads,
+        "num_math_threads_per_group": num_math_threads_per_group,
+    }
+    return hash(frozenset(key_params.items()))
+
+
+class KernelLaunchCache:
+    def __init__(self):
+        self.config_cache = {}
+        self.attr_cache = {}
+
+    @staticmethod
+    def create_attr(kwargs):
+        """Creates and caches property objects"""
+        attr_val = cbd.CUlaunchAttributeValue()
+        attr_val.clusterDim.x = kwargs["NUM_TMA_MULTICAST"]
+        attr_val.clusterDim.y = 1
+        attr_val.clusterDim.z = 1
+        attr = cbd.CUlaunchAttribute()
+        attr.id = cbd.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION
+        attr.value = attr_val
+        return attr
+
+    @staticmethod
+    def create_config(
+        kwargs, num_tma_threads, num_math_threads_per_group, attr
+    ):
+        """Creates and caches configuration objects"""
+        config = cbd.CUlaunchConfig()
+        config.numAttrs = 1
+        config.attrs = [attr]
+        config.gridDimX = kwargs["NUM_SMS"]
+        config.gridDimY = 1
+        config.gridDimZ = 1
+        config.blockDimX = get_num_threads_per_sm(
+            num_tma_threads, num_math_threads_per_group, kwargs["BLOCK_M"]
+        )
+        config.blockDimY = 1
+        config.blockDimZ = 1
+        config.sharedMemBytes = kwargs["SMEM_SIZE"]
+        config.hStream = kwargs["STREAM"]
+        return config
+
+    def get_launch_config(
+        self, kwargs, num_tma_threads, num_math_threads_per_group
+    ):
+        """Retrieves cached config or creates new instance"""
+        cache_key = get_cache_key(
+            kwargs, num_tma_threads, num_math_threads_per_group
+        )
+
+        if cache_key not in self.config_cache:
+            # 先检查属性缓存
+            attr_key = (kwargs["NUM_TMA_MULTICAST"],)
+            if attr_key not in self.attr_cache:
+                self.attr_cache[attr_key] = self.create_attr(kwargs)
+
+            # 创建新配置
+            config = self.create_config(
+                kwargs,
+                num_tma_threads,
+                num_math_threads_per_group,
+                self.attr_cache[attr_key],
+            )
+            self.config_cache[cache_key] = config
+
+        return self.config_cache[cache_key]
+
+
+launch_cache = KernelLaunchCache()
+
+
+class FP8GemmRuntime(Runtime):
+    def __init__(self, path: str) -> None:
+        super().__init__(path)
+
+    @staticmethod
+    def generate(kwargs: dict[str, Any]) -> str:
+        code = f"""
+#ifdef __CUDACC_RTC__
+#include <deep_gemm/nvrtc_std.cuh>
+#else
+#include <cuda.h>
+#include <string>
+#endif
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#include <deep_gemm/fp8_gemm.cuh>
+
+using namespace deep_gemm;
+
+static void __instantiate_kernel() {{
+    auto ptr = reinterpret_cast<void*>(&fp8_gemm_kernel<
+        {kwargs['N']},
+        {kwargs['K']},
+        {kwargs['BLOCK_M']},
+        {kwargs['BLOCK_N']},
+        {kwargs['BLOCK_K']},
+        {kwargs['BLOCK_N_PADDING']},
+        {kwargs['SWIZZLE_D_MODE']},
+        {kwargs['NUM_GROUPS']},
+        {kwargs['NUM_STAGES']},
+        {kwargs['NUM_TMA_THREADS']},
+        {kwargs['NUM_MATH_THREADS_PER_GROUP']},
+        {kwargs['NUM_TMA_MULTICAST']},
+        {'true' if kwargs['IS_TMA_MULTICAST_ON_A'] else 'false'},
+        GemmType::{kwargs['GEMM_TYPE']}
+      >);
+}};
+"""
+        if int(os.getenv("DG_JIT_DEBUG", 0)):
+            print(f"Generated FP8 GEMM code:\n{code}")
+        return code
+
+    # noinspection PyMethodOverriding
+    @staticmethod
+    def launch(kernel: cbd.CUkernel, kwargs: dict[str, Any]) -> cbd.CUresult:
+        num_tma_threads = 128
+        num_math_threads_per_group = 128
+
+        result = cbd.cuKernelSetAttribute(
+            cbd.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
+            kwargs["SMEM_SIZE"],
+            kernel,
+            cbd.CUdevice(kwargs["DEVICE_INDEX"]),
+        )[0]
+        assert result == cbd.CUresult.CUDA_SUCCESS, (
+            f"Failed to set max dynamic shared memory size: {result}"
+        )
+        config = launch_cache.get_launch_config(
+            kwargs, num_tma_threads, num_math_threads_per_group
+        )
+
+        arg_values = (
+            kwargs["SCALES_B"].data_ptr(),
+            kwargs["GROUPED_LAYOUT"].data_ptr(),
+            kwargs["M"],
+            kwargs["TENSOR_MAP_A"],
+            kwargs["TENSOR_MAP_B"],
+            kwargs["TENSOR_MAP_SCALES_A"],
+            kwargs["TENSOR_MAP_D"],
+        )
+        arg_types = (
+            ctypes.c_void_p,
+            ctypes.c_void_p,
+            ctypes.c_uint32,
+            None,
+            None,
+            None,
+            None,
+        )
+        ret = cbd.cuLaunchKernelEx(config, kernel, (arg_values, arg_types), 0)
+        return ret
+
+
+class FP8WGradGemmRuntime(Runtime):
+    def __init__(self, path: str) -> None:
+        super().__init__(path)
+
+    @staticmethod
+    def generate(kwargs: dict[str, Any]) -> str:
+        code = f"""
+#ifdef __CUDACC_RTC__
+#include <deep_gemm/nvrtc_std.cuh>
+#else
+#include <cuda.h>
+#include <string>
+#endif
+
+#include <cuda_bf16.h>
+#include <cuda_fp8.h>
+
+#include <deep_gemm/fp8_wgrad_gemm.cuh>
+
+using namespace deep_gemm;
+
+static void __instantiate_kernel() {{
+    auto ptr = reinterpret_cast<void*>(&fp8_wgrad_gemm_kernel<
+        {kwargs['M']},
+        {kwargs['N']},
+        {kwargs['BLOCK_M']},
+        {kwargs['BLOCK_N']},
+        {kwargs['BLOCK_K']},
+        {kwargs['NUM_STAGES']},
+        {kwargs['NUM_LAST_STAGES']},
+        {kwargs['NUM_TMA_THREADS']},
+        {kwargs['NUM_MATH_THREADS_PER_GROUP']},
+        {kwargs['NUM_TMA_MULTICAST']},
+        {'true' if kwargs['IS_TMA_MULTICAST_ON_A'] else 'false'}
+      >);
+}};
+"""
+        if int(os.getenv("DG_JIT_DEBUG", 0)):
+            print(f"Generated FP8 WGrad GEMM code:\n{code}")
+        return code
+
+    # noinspection PyMethodOverriding
+    @staticmethod
+    def launch(kernel: cbd.CUkernel, kwargs: dict[str, Any]) -> cbd.CUresult:
+        num_tma_threads = 128
+        num_math_threads_per_group = 128
+
+        result = cbd.cuKernelSetAttribute(
+            cbd.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
+            kwargs["SMEM_SIZE"],
+            kernel,
+            cbd.CUdevice(kwargs["DEVICE_INDEX"]),
+        )[0]
+        assert result == cbd.CUresult.CUDA_SUCCESS, (
+            f"Failed to set max dynamic shared memory size: {result}"
+        )
+        config = launch_cache.get_launch_config(
+            kwargs, num_tma_threads, num_math_threads_per_group
+        )
+
+        arg_values = (
+            kwargs["K"],
+            kwargs["TENSOR_MAP_A"],
+            kwargs["TENSOR_MAP_B"],
+            kwargs["TENSOR_MAP_SCALES_A"],
+            kwargs["TENSOR_MAP_SCALES_B"],
+            kwargs["TENSOR_MAP_D"],
+        )
+        arg_types = (
+            ctypes.c_uint32,
+            None,
+            None,
+            None,
+            None,
+            None,
+        )
+        return cbd.cuLaunchKernelEx(config, kernel, (arg_values, arg_types), 0)

python/paddle/incubate/fp8/deep_gemm/jit_kernels/__init__.py

@@ -0,0 +1,34 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from .gemm import gemm_fp8_fp8_bf16_nt  # noqa: F401
+from .m_grouped_gemm import (  # noqa: F401
+    m_grouped_gemm_fp8_fp8_bf16_nt_contiguous,
+    m_grouped_gemm_fp8_fp8_bf16_nt_masked,
+)
+from .utils import (  # noqa: F401
+    ceil_div,
+    get_col_major_tma_aligned_tensor,
+    get_m_alignment_for_contiguous_layout,
+    get_num_sms,
+    set_num_sms,
+)
+from .wgrad_gemm import (  # noqa: F401
+    k_grouped_wgrad_gemm_fp8_fp8_fp32_nt,
+    wgrad_gemm_fp8_fp8_fp32_nt,
+)

python/paddle/incubate/fp8/deep_gemm/jit_kernels/gemm.py

@@ -0,0 +1,394 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from __future__ import annotations
+
+import math
+import os
+from functools import cache
+
+import paddle
+
+from ..jit import FP8GemmRuntime, build
+from .runtime import (
+    GemmType,
+    make_2d_tma_a_desc,
+    make_2d_tma_b_desc,
+    make_2d_tma_d_desc,
+    make_2d_tma_scales_desc,
+)
+from .utils import (
+    ceil_div,
+    get_col_major_tma_aligned_tensor,
+    get_m_alignment_for_contiguous_layout,
+    get_num_sms,
+)
+
+global_empty_tensor = paddle.empty([0], dtype=paddle.int32)
+# Todo: Use default stream to accelerate CPU time. Optimize here if use multistream to launch gemm kernel.
+global_stream = paddle.device.current_stream().stream_base.cuda_stream
+
+
+def is_tma_multicast_legal(
+    shape_dim: int,
+    block_dim: int,
+    num_tma_multicast: int,
+    num_sms: int,
+    require_divisible: bool = False,
+) -> bool:
+    divisible = (
+        ceil_div(shape_dim, block_dim) % num_tma_multicast == 0
+        or not require_divisible
+    )
+    return divisible and num_sms % num_tma_multicast == 0
+
+
+def get_swizzle_mode(block_n: int) -> int:
+    elem_size = 2
+    for mode_bytes in (128, 64, 32):
+        if (block_n * elem_size) % mode_bytes == 0:
+            return mode_bytes
+    return 0
+
+
+def get_block_n_padding_for_smem_d(block_n: int) -> int:
+    # NOTES: padding is for solving bank conflicts, but wastes shared memory space
+    elem_size, requirement = 2, (4, 8)
+    bank_stride = (block_n * elem_size) // 4
+    padding = (requirement[0] - bank_stride) % requirement[1]
+    return (
+        ((padding + requirement[1]) if padding < 0 else padding) * 4
+    ) // elem_size
+
+
+def get_smem_config(
+    num_stages: int,
+    k: int,
+    block_m: int,
+    block_n: int,
+    block_k: int = 128,
+    is_fp32_out: bool = False,
+    is_wgrad: bool = False,
+) -> tuple[int, int, int]:
+    assert block_k == 128
+
+    # Try swizzle first, as it does not waste shared memory
+    swizzle_mode = get_swizzle_mode(block_n)
+    block_n_padding = (
+        get_block_n_padding_for_smem_d(block_n) if swizzle_mode == 0 else 0
+    )
+
+    # NOTES: `scales_b` in a total manner or per-stage manner
+    smem_d = block_m * (block_n + block_n_padding) * (4 if is_fp32_out else 2)
+    smem_a_per_stage = block_m * block_k
+    smem_scales_a_per_stage = block_m * 4
+    smem_b_per_stage = block_n * block_k
+    smem_scales_b_per_stage = (
+        ceil_div(block_n * 4, block_k) * block_k if is_wgrad else 0
+    )
+    smem_scales_b = ceil_div(k, block_k) * 4 if not is_wgrad else 0
+    smem_barrier = num_stages * 8 * 2
+
+    smem_size = 0
+    smem_size += smem_d
+    smem_size += num_stages * smem_a_per_stage
+    smem_size += num_stages * smem_scales_a_per_stage
+    smem_size += num_stages * smem_b_per_stage
+    smem_size += num_stages * smem_scales_b_per_stage
+    smem_size += (
+        ceil_div(smem_scales_b * (1 if block_k % block_n == 0 else 2), 8) * 8
+    )
+    smem_size += smem_barrier
+
+    # Swizzle and padding are not compatible
+    assert int(swizzle_mode > 0) + int(block_n_padding > 0) <= 1
+
+    return smem_size, swizzle_mode, block_n_padding
+
+
+@cache
+def get_best_configs(
+    m: int,
+    n: int,
+    k: int,
+    num_groups: int,
+    num_sms: int,
+    is_grouped_contiguous: bool = False,
+    is_grouped_masked: bool = False,
+    is_fp32_out: bool = False,
+    is_wgrad: bool = False,
+) -> tuple[int, int, int, int, tuple[int, bool], tuple[int, int, int]]:
+    if not is_grouped_contiguous:
+        block_ms = (
+            64,
+            128,
+        ) + ((256,) if not is_fp32_out else ())
+    else:
+        block_ms = (get_m_alignment_for_contiguous_layout(),)
+    block_ns = tuple(range(16, 129, 8)) + (
+        (
+            136,
+            152,
+        )
+        if is_wgrad
+        else (
+            144,
+            160,
+        )
+    )
+
+    # Avoid bank conflicts for FP32 output
+    if is_fp32_out:
+        block_ns = [x for x in block_ns if x % 16 == 8]
+
+    fix_wave_saturate = lambda x: num_sms if x == 0 else x
+    get_num_waves = lambda bm, bn: (
+        ceil_div(ceil_div(m, bm) * ceil_div(n, bn) * num_groups, num_sms)
+        if bm
+        else None
+    )
+    get_last_wave_util = lambda bm, bn: fix_wave_saturate(
+        (ceil_div(m, bm) * ceil_div(n, bn) * num_groups) % num_sms
+    )
+
+    # Decide block sizes by waves
+    best_block_m, best_block_n = None, None
+    for block_m in block_ms:
+        # NOTES: the block sizes cannot be too large, so at least one dim less than 128
+        for block_n in filter(lambda bn: block_m <= 128 or bn <= 128, block_ns):
+            success = False
+            num_waves, best_num_waves = (
+                get_num_waves(block_m, block_n),
+                get_num_waves(best_block_m, best_block_n),
+            )
+            if best_block_m is None or best_block_n is None:
+                success = True
+            elif num_waves < best_num_waves:
+                success = True
+            elif num_waves == best_num_waves:
+                # Check last wave utilization
+                util = get_last_wave_util(block_m, block_n)
+                best_util = get_last_wave_util(best_block_m, best_block_n)
+                success = util > best_util
+                if util == best_util:
+                    # Case 1: same `block_m`, smaller `block_n` (wasted)
+                    success |= (
+                        block_m == best_block_m and block_n < best_block_n
+                    )
+                    # Case 2: same `block_n`, smaller `block_m` (wasted)
+                    success |= (
+                        block_n == best_block_n and block_m < best_block_m
+                    )
+                    # Case 3: different for both `block_m` and `block_n`, `block_n` larger is better
+                    success |= (
+                        block_m != best_block_m and block_n > best_block_n
+                    )
+            best_block_m, best_block_n = (
+                (block_m, block_n) if success else (best_block_m, best_block_n)
+            )
+    assert best_block_m is not None and best_block_n is not None
+
+    # Always pick the longest one
+    # NOTES: for double B scales, the best number of stages may be reduced
+    best_num_stages, best_smem_config, sm90_capacity = None, None, 232448
+    stage_candidates = tuple(
+        filter(lambda s: s <= max(k // 128, 1), (8, 7, 6, 5, 4, 3, 2, 1))
+    )
+    if 128 % best_block_n != 0 and 128 // math.gcd(128, best_block_n) <= 4:
+        # Unrolling both stages and `num_former_iters` will cause large code size
+        stage_candidates = tuple(
+            filter(lambda s: s <= max(k // 128, 1), (4, 3, 2, 1))
+        )
+    for num_stages in stage_candidates:
+        best_smem_config = get_smem_config(
+            num_stages,
+            k,
+            best_block_m,
+            best_block_n,
+            is_fp32_out=is_fp32_out,
+            is_wgrad=is_wgrad,
+        )
+        if best_smem_config[0] <= sm90_capacity:
+            best_num_stages = num_stages
+            break
+    assert best_smem_config is not None
+    assert best_num_stages is not None
+
+    # Decide the number of TMA multicasts and whether broadcast on A
+    best_tma_multicast_config = (1, True)
+
+    # Try to multicast on the larger block side first
+    # NOTES: currently, grouped masked GEMM only supports multicast on A and requires the number of blocks in the N-direction to be even
+    is_multicast_legal = {
+        "A": is_tma_multicast_legal(
+            n, best_block_n, 2, num_sms, is_grouped_masked
+        ),
+        "B": is_tma_multicast_legal(m, best_block_m, 2, num_sms)
+        and not is_grouped_masked,
+    }
+    for i in ("A", "B") if best_block_m > best_block_n else ("B", "A"):
+        if m >= 512 and is_multicast_legal[i]:
+            best_tma_multicast_config = (2, i == "A")
+            break
+
+    # Recompute the minimal number of SMs required
+    # NOTES: less L2 cache usage and less GPU frequency drop
+    num_waves = get_num_waves(best_block_m, best_block_n)
+    num_min_sms = ceil_div(
+        ceil_div(m, best_block_m) * ceil_div(n, best_block_n) * num_groups,
+        num_waves,
+    )
+    num_min_sms = (
+        ceil_div(num_min_sms, best_tma_multicast_config[0])
+        * best_tma_multicast_config[0]
+    )
+    assert num_min_sms <= num_sms
+
+    return (
+        num_min_sms,
+        best_block_m,
+        best_block_n,
+        best_num_stages,
+        best_tma_multicast_config,
+        best_smem_config,
+    )
+
+
+def gemm_fp8_fp8_bf16_nt(
+    lhs: tuple[paddle.Tensor, paddle.Tensor],
+    rhs: tuple[paddle.Tensor, paddle.Tensor],
+    out: paddle.Tensor,
+    num_sms: int | None = None,
+) -> None:
+    """
+    Perform a normal GEMM with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
+
+    Requirements:
+        LHS, RHS, and output tensors must be contiguous in dimension 1, i.e., strides[1] = 1.
+        The strides[0] of LHS and RHS must be a multiple of 16, and the strides[0] of output must be a multiple of 8.
+        RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
+            this function will do a transposing with a set of slow PaddlePaddle operations.
+
+    Arguments:
+        lhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[m, k]`,
+             the second element is an FP32 1x128 scaling tensor for LHS of shape `[m, ⌈k / 128⌉]`.
+        rhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[n, k]`,
+             the second element is an FP32 128x128 scaling tensor for RHS of shape `[⌈n / 128⌉, ⌈k / 128⌉]`.
+        out: the BF16 output tensor of shape `[m, n]`, representing the result.
+    """
+    lhs, lhs_scales = lhs
+    rhs, rhs_scales = rhs
+    m, k = lhs.shape
+    n, k_ = rhs.shape
+    m_, n_ = out.shape
+
+    # Type and shape checks
+    assert m == m_ and n == n_ and k == k_
+    assert n > 0 and k > 0
+    assert lhs_scales.shape == [m, ceil_div(k, 128)]
+    assert rhs_scales.shape == [ceil_div(n, 128), ceil_div(k, 128)]
+    assert (
+        lhs.dtype == paddle.float8_e4m3fn and lhs_scales.dtype == paddle.float32
+    )
+    assert (
+        rhs.dtype == paddle.float8_e4m3fn and rhs_scales.dtype == paddle.float32
+    )
+    assert out.dtype == paddle.bfloat16
+    assert lhs.strides[1] == 1 and out.strides[1] == 1 and rhs.strides[1] == 1
+
+    # LHS scales must be transposed for TMA loads, but not for RHS scales
+    # NOTES: `get_col_major_tma_aligned_tensor` may launch a kernel if not processed by previous kernels
+    lhs_scales = get_col_major_tma_aligned_tensor(lhs_scales)
+    assert rhs_scales.is_contiguous()
+
+    # Do nothing if `m` is zero
+    if m == 0:
+        return
+
+    # K must be aligned to 128
+    aligned_k = ceil_div(k, 128) * 128
+
+    # Auto-tuning with compilation
+    if num_sms is None:
+        num_sms = get_num_sms()
+    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = (
+        get_best_configs(m, n, k, 1, num_sms)
+    )
+    if int(os.getenv("DG_JIT_KERNELS_DEBUG", 0)):
+        print(
+            f"Auto-tuned gemm_fp8_fp8_bf16_nt as num_sms={num_sms}, block_m={block_m}, block_n={block_n}"
+        )
+    block_k = 128
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    tensor_map_a = make_2d_tma_a_desc(
+        GemmType.Normal, lhs, m, k, lhs.strides[0], block_m, block_k, 1
+    )
+    tensor_map_b = make_2d_tma_b_desc(
+        GemmType.Normal, rhs, n, k, rhs.strides[0], block_n, block_k, 1
+    )
+    tensor_map_d = make_2d_tma_d_desc(
+        GemmType.Normal,
+        out,
+        m,
+        n,
+        out.strides[0],
+        block_m,
+        block_n,
+        1,
+        smem_config[1],
+    )
+    tensor_map_scales_a = make_2d_tma_scales_desc(
+        GemmType.Normal, lhs_scales, m, k, block_m, block_k, 1
+    )
+
+    kwargs = {
+        # Templated arguments
+        "GEMM_TYPE": GemmType.Normal,
+        "NUM_TMA_THREADS": num_tma_threads,
+        "NUM_MATH_THREADS_PER_GROUP": num_math_threads_per_group,
+        "M": m,
+        "N": n,
+        "K": aligned_k,
+        "NUM_GROUPS": 1,
+        "BLOCK_M": block_m,
+        "BLOCK_N": block_n,
+        "BLOCK_K": block_k,
+        "SWIZZLE_D_MODE": smem_config[1],
+        "BLOCK_N_PADDING": smem_config[2],
+        "NUM_STAGES": num_stages,
+        "NUM_TMA_MULTICAST": tma_multicast_config[0],
+        "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+        # Runtime arguments
+        "SCALES_B": rhs_scales,
+        "GROUPED_LAYOUT": global_empty_tensor,
+        "NUM_SMS": num_sms,
+        "SMEM_SIZE": smem_config[0],
+        "TENSOR_MAP_A": tensor_map_a,
+        "TENSOR_MAP_B": tensor_map_b,
+        "TENSOR_MAP_SCALES_A": tensor_map_scales_a,
+        "TENSOR_MAP_D": tensor_map_d,
+        "STREAM": global_stream,
+        "DEVICE_INDEX": out.place.gpu_device_id(),
+    }
+
+    # Generate, build and run the kernel
+    runtime = build("gemm_fp8_fp8_bf16_nt", FP8GemmRuntime, kwargs)
+    runtime(**kwargs)

python/paddle/incubate/fp8/deep_gemm/jit_kernels/m_grouped_gemm.py

@@ -0,0 +1,310 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from __future__ import annotations
+
+import os
+from functools import reduce
+
+import paddle
+
+from ..jit import FP8GemmRuntime, build
+from .gemm import get_best_configs
+from .runtime import (
+    GemmType,
+    make_2d_tma_a_desc,
+    make_2d_tma_b_desc,
+    make_2d_tma_d_desc,
+    make_2d_tma_scales_desc,
+)
+from .utils import ceil_div, get_col_major_tma_aligned_tensor, get_num_sms
+
+# Todo: Use default stream to accelerate CPU time. Optimize here if use multistream to launch gemm kernel.
+global_stream = paddle.device.current_stream().stream_base.cuda_stream
+
+
+def m_grouped_gemm_fp8_fp8_bf16_nt_contiguous(
+    lhs: tuple[paddle.Tensor, paddle.Tensor],
+    rhs: tuple[paddle.Tensor, paddle.Tensor],
+    out: paddle.Tensor,
+    m_indices: paddle.Tensor,
+    num_sms: int | None = None,
+) -> None:
+    """
+    Perform a grouped GEMM (contiguous format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
+
+    Requirements:
+        LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
+        RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
+            this function will do a transposing with a set of slow PaddlePaddle operations.
+        On the M axis, inputs are grouped into several batches, of which batch sizes aligned to
+            `get_m_alignment_for_contiguous_layout()` (128).
+
+    Arguments:
+        lhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[m_sum, k]`,
+             the second element is an FP32 1x128 scaling tensor for LHS of shape `[m_sum, ⌈k / 128⌉]`.
+        rhs: the first element is an FP8 tensor (typed `paddle.float8_e4m3fn`) of shape `[num_groups, n, k]`,
+             the second element is an FP32 128x128 scaling tensor for RHS of shape `[num_groups, ⌈n / 128⌉, ⌈k / 128⌉]`.
+        out: the BF16 output tensor of shape `[m_sum, n]`, representing the result.
+        m_indices: a tensor of shape `[m_sum]` with type `paddle.int`.
+            `m_indices[i]` records the group which the i-th row of the LHS belongs to,
+            which means that the i-th row of the LHS matrix will be multiplied with `rhs[m_indices[i]]`.
+            Values of `m_indices` in every-m-alignment-block must also be the same.
+    """
+    lhs, lhs_scales = lhs
+    rhs, rhs_scales = rhs
+    m, k = lhs.shape
+    num_groups, n, k_ = rhs.shape
+    m_, n_ = out.shape
+    m_shape = m_indices.shape
+    m__ = reduce(lambda x, y: x * y, m_shape)
+
+    # Type and shape checks
+    assert m == m_ == m__ and k == k_ and n == n_
+    assert lhs_scales.shape == [m, ceil_div(k, 128)]
+    assert rhs_scales.shape == [num_groups, ceil_div(n, 128), ceil_div(k, 128)]
+    assert (
+        lhs.dtype == paddle.float8_e4m3fn and lhs_scales.dtype == paddle.float32
+    )
+    assert (
+        rhs.dtype == paddle.float8_e4m3fn and rhs_scales.dtype == paddle.float32
+    )
+    assert out.dtype == paddle.bfloat16
+    assert m_indices.dtype == paddle.int32
+    assert lhs.is_contiguous() and rhs.is_contiguous()
+    assert out.is_contiguous() and m_indices.is_contiguous()
+
+    # LHS scales must be transposed for TMA load, but not for RHS scales
+    lhs_scales = get_col_major_tma_aligned_tensor(lhs_scales)
+    assert rhs_scales.is_contiguous()
+
+    # Do nothing if `m` is zero
+    if m == 0:
+        return
+
+    # Auto-tuning with compilation
+    if num_sms is None:
+        num_sms = get_num_sms()
+    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = (
+        get_best_configs(m, n, k, 1, num_sms, is_grouped_contiguous=True)
+    )
+    if int(os.getenv("DG_JIT_KERNELS_DEBUG", 0)):
+        print(
+            f"Auto-tuned m_grouped_gemm_fp8_fp8_bf16_nt_contiguous as num_sms={num_sms}, block_m={block_m}, block_n={block_n}"
+        )
+    block_k = 128
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    tensor_map_a = make_2d_tma_a_desc(
+        GemmType.GroupedContiguous, lhs, m, k, k, block_m, block_k, num_groups
+    )
+    tensor_map_b = make_2d_tma_b_desc(
+        GemmType.GroupedContiguous, rhs, n, k, k, block_n, block_k, num_groups
+    )
+    tensor_map_d = make_2d_tma_d_desc(
+        GemmType.GroupedContiguous,
+        out,
+        m,
+        n,
+        n,
+        block_m,
+        block_n,
+        num_groups,
+        smem_config[1],
+    )
+    tensor_map_scales_a = make_2d_tma_scales_desc(
+        GemmType.GroupedContiguous,
+        lhs_scales,
+        m,
+        k,
+        block_m,
+        block_k,
+        num_groups,
+    )
+
+    kwargs = {
+        # Templated arguments
+        "NUM_TMA_THREADS": num_tma_threads,
+        "NUM_MATH_THREADS_PER_GROUP": num_math_threads_per_group,
+        "M": m,
+        "N": n,
+        "K": k,
+        "BLOCK_M": block_m,
+        "BLOCK_N": block_n,
+        "BLOCK_K": block_k,
+        "SWIZZLE_D_MODE": smem_config[1],
+        "BLOCK_N_PADDING": smem_config[2],
+        "NUM_GROUPS": num_groups,
+        "NUM_STAGES": num_stages,
+        "NUM_TMA_MULTICAST": tma_multicast_config[0],
+        "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+        "GEMM_TYPE": GemmType.GroupedContiguous,
+        # Runtime arguments
+        "SCALES_B": rhs_scales,
+        "GROUPED_LAYOUT": m_indices,
+        "NUM_SMS": num_sms,
+        "SMEM_SIZE": smem_config[0],
+        "TENSOR_MAP_A": tensor_map_a,
+        "TENSOR_MAP_B": tensor_map_b,
+        "TENSOR_MAP_SCALES_A": tensor_map_scales_a,
+        "TENSOR_MAP_D": tensor_map_d,
+        "STREAM": global_stream,
+        "DEVICE_INDEX": out.place.gpu_device_id(),
+    }
+
+    # Generate, build and run the kernel
+    runtime = build("m_grouped_gemm_fp8_fp8_bf16_nt", FP8GemmRuntime, kwargs)
+    runtime(**kwargs)
+
+
+def m_grouped_gemm_fp8_fp8_bf16_nt_masked(
+    lhs: tuple[paddle.Tensor, paddle.Tensor],
+    rhs: tuple[paddle.Tensor, paddle.Tensor],
+    out: paddle.Tensor,
+    masked_m: paddle.Tensor,
+    expected_m: int,
+    num_sms: int | None = None,
+) -> None:
+    """
+    Perform a grouped GEMM (masked format) with FP8 inputs and BF16 output, with 1x128 LHS scaling and 128x128 RHS scaling.
+
+    Requirements:
+        LHS, RHS, RHS scaling factors, and output tensors must be in contiguous format.
+        RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling tensor requires a TMA-aligned transposed format, if your input does not match the requirement,
+            this function will do a transposing with a set of slow PaddlePaddle operations.
+        Moreover, this alignment requirement is different with the contiguous-format kernel, as we require that each batch
+            should be separately transposed.
+
+    Arguments:
+        lhs: the first element is an FP8 tensor (typed `paddle.bfloat16`) of shape `[num_groups, m_max, k]`,
+             the second element is an FP32 1x128 scaling tensor for LHS of shape `[num_groups, m_max, ⌈k / 128⌉]`.
+        rhs: the first element is an FP8 tensor (typed `paddle.bfloat16`) of shape `[num_groups, n, k]`.
+             The second element is an FP32 128x128 scaling tensor for RHS of shape `[num_groups, ⌈n / 128⌉, ⌈k / 128⌉]`.
+        out: the BF16 output tensor of shape `[num_groups, m_max, n]`, representing the result.
+        masked_m: a tensor of shape `[num_groups]`, `masked_m[i]` records actual rows of the `lhs[i]` matrix to compute
+            in the i-th group.
+        expected_m: a value hint (which is a value on CPU) for the M expectation of each batch,
+            correctly setting this value may lead to better performance.
+    """
+    lhs, lhs_scales = lhs
+    rhs, rhs_scales = rhs
+    num_groups, m, k = lhs.shape
+    num_groups_, n, k_ = rhs.shape
+    num_groups__, m_, n_ = out.shape
+    num_groups___ = masked_m.shape[0]
+
+    # Type and shape checks
+    assert num_groups == num_groups_ == num_groups__ == num_groups___
+    assert m == m_ and n == n_ and k == k_
+    assert expected_m > 0 and m > 0 and n > 0 and k > 0 and num_groups > 0
+    assert lhs_scales.shape == [num_groups, m, ceil_div(k, 128)]
+    assert rhs_scales.shape == [num_groups, ceil_div(n, 128), ceil_div(k, 128)]
+    assert (
+        lhs.dtype == paddle.float8_e4m3fn and lhs_scales.dtype == paddle.float32
+    )
+    assert (
+        rhs.dtype == paddle.float8_e4m3fn and rhs_scales.dtype == paddle.float32
+    )
+    assert out.dtype == paddle.bfloat16
+    assert masked_m.dtype == paddle.int32
+    assert lhs.is_contiguous() and rhs.is_contiguous()
+    assert out.is_contiguous() and masked_m.is_contiguous()
+
+    # LHS scales must be transposed for TMA load, but not for RHS scales
+    lhs_scales = get_col_major_tma_aligned_tensor(lhs_scales)
+    assert rhs_scales.is_contiguous()
+
+    # Auto-tuning with compilation
+    if num_sms is None:
+        num_sms = get_num_sms()
+    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = (
+        get_best_configs(
+            expected_m, n, k, num_groups, num_sms, is_grouped_masked=True
+        )
+    )
+    if int(os.getenv("DG_JIT_KERNELS_DEBUG", 0)):
+        print(
+            f"Auto-tuned m_grouped_gemm_fp8_fp8_bf16_nt_masked as num_sms={num_sms}, block_m={block_m}, block_n={block_n}"
+        )
+    # Extra checks for TMA store
+    if num_groups > 1 and m > block_m:
+        assert m % block_m == 0, (
+            f"For masked grouped GEMM, shape M should be multiple of the block M (current block M: {block_m})"
+        )
+
+    block_k = 128
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    tensor_map_a = make_2d_tma_a_desc(
+        GemmType.GroupedMasked, lhs, m, k, k, block_m, block_k, num_groups
+    )
+    tensor_map_b = make_2d_tma_b_desc(
+        GemmType.GroupedMasked, rhs, n, k, k, block_n, block_k, num_groups
+    )
+    tensor_map_d = make_2d_tma_d_desc(
+        GemmType.GroupedMasked,
+        out,
+        m,
+        n,
+        n,
+        block_m,
+        block_n,
+        num_groups,
+        smem_config[1],
+    )
+    tensor_map_scales_a = make_2d_tma_scales_desc(
+        GemmType.GroupedMasked, lhs_scales, m, k, block_m, block_k, num_groups
+    )
+
+    kwargs = {
+        # Templated arguments
+        "NUM_TMA_THREADS": num_tma_threads,
+        "NUM_MATH_THREADS_PER_GROUP": num_math_threads_per_group,
+        "M": m,
+        "N": n,
+        "K": k,
+        "BLOCK_M": block_m,
+        "BLOCK_N": block_n,
+        "BLOCK_K": block_k,
+        "SWIZZLE_D_MODE": smem_config[1],
+        "BLOCK_N_PADDING": smem_config[2],
+        "NUM_GROUPS": num_groups,
+        "NUM_STAGES": num_stages,
+        "NUM_TMA_MULTICAST": tma_multicast_config[0],
+        "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+        "GEMM_TYPE": GemmType.GroupedMasked,
+        # Runtime arguments
+        "SCALES_B": rhs_scales,
+        "GROUPED_LAYOUT": masked_m,
+        "NUM_SMS": num_sms,
+        "SMEM_SIZE": smem_config[0],
+        "TENSOR_MAP_A": tensor_map_a,
+        "TENSOR_MAP_B": tensor_map_b,
+        "TENSOR_MAP_SCALES_A": tensor_map_scales_a,
+        "TENSOR_MAP_D": tensor_map_d,
+        "STREAM": paddle.device.cuda.current_stream().cuda_stream,
+        "DEVICE_INDEX": out.place.gpu_device_id(),
+    }
+
+    # Generate, build and run the kernel
+    runtime = build("m_grouped_gemm_fp8_fp8_bf16_nt", FP8GemmRuntime, kwargs)
+    runtime(**kwargs)

python/paddle/incubate/fp8/deep_gemm/jit_kernels/runtime.py

@@ -0,0 +1,186 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from typing import Any
+
+import cuda.bindings.driver as cbd
+
+import paddle
+
+from ..jit.runtime import GemmType
+from .utils import get_tma_aligned_size
+
+# TODO Support dtype in Paddle
+tmap_type_map: dict[Any, str] = {
+    paddle.int8: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    paddle.int16: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT16,
+    paddle.int32: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_INT32,
+    paddle.int64: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_INT64,
+    paddle.uint8: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    # paddle.uint16:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT16,
+    # paddle.uint32:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT32,
+    # paddle.uint64:          cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT64,
+    paddle.float32: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_FLOAT32,
+    paddle.float16: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_FLOAT16,
+    paddle.bfloat16: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_BFLOAT16,
+    paddle.float8_e4m3fn: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    # paddle.float8_e4m3fnuz: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    paddle.float8_e5m2: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+    # paddle.float8_e5m2fnuz: cbd.CUtensorMapDataType.CU_TENSOR_MAP_DATA_TYPE_UINT8,
+}
+
+swizzle_type_map = {
+    0: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE,
+    32: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_32B,
+    64: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_64B,
+    128: cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B,
+}
+
+
+def make_2d_tma_copy_desc(
+    t: paddle.Tensor,
+    gmem_dims: tuple[cbd.cuuint64_t, cbd.cuuint64_t],
+    gmem_outer_stride: cbd.cuuint64_t,
+    smem_dims: tuple[cbd.cuuint32_t, cbd.cuuint32_t],
+    swizzle_type: cbd.CUtensorMapSwizzle,
+) -> cbd.CUtensorMap:
+    tensor_dtype = tmap_type_map[t.dtype]
+    res, tensor_map = cbd.cuTensorMapEncodeTiled(
+        tensor_dtype,
+        2,
+        t.data_ptr(),
+        gmem_dims,
+        (gmem_outer_stride,),
+        smem_dims,
+        (cbd.cuuint32_t(1), cbd.cuuint32_t(1)),
+        cbd.CUtensorMapInterleave.CU_TENSOR_MAP_INTERLEAVE_NONE,
+        swizzle_type,
+        cbd.CUtensorMapL2promotion.CU_TENSOR_MAP_L2_PROMOTION_L2_256B,
+        cbd.CUtensorMapFloatOOBfill.CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE,
+    )
+    if res != cbd.CUresult.CUDA_SUCCESS:
+        raise Exception(f"Failed to encode tensor map: {res}")
+    return tensor_map
+
+
+def make_2d_tma_desc(
+    t: paddle.Tensor,
+    gmem_inner_dim: int,
+    gmem_outer_dim: int,
+    gmem_outer_stride: int,
+    smem_inner_dim: int,
+    smem_outer_dim: int,
+    swizzle_type: cbd.CUtensorMapSwizzle = cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_128B,
+) -> cbd.CUtensorMap:
+    gmem_dim = (cbd.cuuint64_t(gmem_inner_dim), cbd.cuuint64_t(gmem_outer_dim))
+    smem_dim = (cbd.cuuint32_t(smem_inner_dim), cbd.cuuint32_t(smem_outer_dim))
+    return make_2d_tma_copy_desc(
+        t,
+        gmem_dim,
+        cbd.cuuint64_t(gmem_outer_stride * t.element_size()),
+        smem_dim,
+        swizzle_type,
+    )
+
+
+def make_2d_tma_a_desc(
+    gemm_type: GemmType,
+    t: paddle.Tensor,
+    shape_m: int,
+    shape_k: int,
+    m_stride: int,
+    block_m: int,
+    block_k: int,
+    num_groups: int,
+) -> cbd.CUtensorMap:
+    return make_2d_tma_desc(
+        t,
+        shape_k,
+        shape_m * (num_groups if gemm_type == GemmType.GroupedMasked else 1),
+        m_stride,
+        block_k,
+        block_m,
+    )
+
+
+def make_2d_tma_b_desc(
+    gemm_type: GemmType,
+    t: paddle.Tensor,
+    shape_n: int,
+    shape_k: int,
+    n_stride: int,
+    block_n: int,
+    block_k: int,
+    num_groups: int,
+) -> cbd.CUtensorMap:
+    return make_2d_tma_desc(
+        t,
+        shape_k,
+        shape_n * (num_groups if gemm_type != GemmType.Normal else 1),
+        n_stride,
+        block_k,
+        block_n,
+    )
+
+
+def make_2d_tma_d_desc(
+    gemm_type: GemmType,
+    t: paddle.Tensor,
+    shape_m: int,
+    shape_n: int,
+    m_stride: int,
+    block_m: int,
+    block_n: int,
+    num_groups: int,
+    swizzle_mode: int,
+) -> cbd.CUtensorMap:
+    # Swizzling requires the inner box dim to be less or equal than `kSwizzleDMode`
+    # bytes, so `BLOCK_N * sizeof(T) / kSwizzleDMode` TMA stores are required
+    return make_2d_tma_desc(
+        t,
+        shape_n,
+        shape_m * (num_groups if gemm_type == GemmType.GroupedMasked else 1),
+        m_stride,
+        block_n if swizzle_mode == 0 else swizzle_mode // t.element_size(),
+        block_m,
+        swizzle_type_map[swizzle_mode],
+    )
+
+
+def make_2d_tma_scales_desc(
+    gemm_type: GemmType,
+    t: paddle.Tensor,
+    shape_mn: int,
+    shape_k: int,
+    block_mn: int,
+    block_k: int,
+    num_groups: int,
+) -> cbd.CUtensorMap:
+    # Make TMA aligned to 16 bytes
+    shape_mn = get_tma_aligned_size(shape_mn, t.element_size())
+    return make_2d_tma_desc(
+        t,
+        shape_mn,
+        (shape_k + block_k - 1)
+        // block_k
+        * (num_groups if gemm_type == GemmType.GroupedMasked else 1),
+        shape_mn,
+        block_mn,
+        1,
+        cbd.CUtensorMapSwizzle.CU_TENSOR_MAP_SWIZZLE_NONE,
+    )

python/paddle/incubate/fp8/deep_gemm/jit_kernels/utils.py

@@ -0,0 +1,146 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+import paddle
+from paddle import Tensor
+
+_num_sms = None
+
+
+def set_num_sms(num_sms: int) -> None:
+    """
+    Set the maximum SM count for all GEMM kernels to use.
+
+    Arguments:
+        num_sms: the desired maximum SM count for all GEMM kernels to use.
+    """
+    global _num_sms
+    assert (
+        0
+        < num_sms
+        <= paddle.device.cuda.get_device_properties(
+            device="cuda"
+        ).multi_processor_count
+    )
+    _num_sms = num_sms
+
+
+def get_num_sms() -> int:
+    """
+    Get the current maximum limit of SM count for all GEMM kernels to use.
+    If the count is never specified, the function will return the number of device SMs.
+
+    Returns:
+        Current maximum limit of SM count for all GEMM kernels to use.
+    """
+    global _num_sms
+    if _num_sms is None:
+        _num_sms = (
+            paddle.device.cuda.get_device_properties().multi_processor_count
+        )
+    return _num_sms
+
+
+def ceil_div(x: int, y: int) -> int:
+    """
+    Perform ceiling division of two integers.
+
+    Args:
+        x: the dividend.
+        y: the divisor.
+
+    Returns:
+        The result of the ceiling division.
+    """
+    return (x + y - 1) // y
+
+
+def get_m_alignment_for_contiguous_layout():
+    """
+    When we do a grouped GEMM in contiguous format, LHS are grouped into several batches along the M axis.
+    Since we deal with exactly one sub-matrix of RHS for each GEMM block, batch sizes above should align well
+        with GEMM block shape.
+
+    Returns:
+        Group-level alignment requirement for grouped contiguous layout, which is always 128.
+    """
+    return 128
+
+
+def get_tma_aligned_size(x: int, element_size: int) -> int:
+    """
+    Global memory address of TMA must be 16-byte aligned.
+    Since we use column-major layout for the LHS scaling tensor,
+        the M-axis of the LHS scaling tensor needs to be padded to a multiple of 16 bytes.
+
+    Arguments:
+        x: original M-axis shape of the LHS scaling tensor.
+        element_size: element size of the LHS scaling tensor.
+
+    Returns:
+        M-axis shape of the LHS scaling tensor after padding.
+    """
+    tma_alignment_bytes = 16
+    assert tma_alignment_bytes % element_size == 0
+    alignment = tma_alignment_bytes // element_size
+    return ceil_div(x, alignment) * alignment
+
+
+def get_col_major_tma_aligned_tensor(x: Tensor) -> Tensor:
+    """
+    Returns TMA-aligned transposed format of the input tensor. `paddle.transpose` will be called if necessary.
+    If the input tensor is already column-major layout and 16-byte aligned along the M axis
+        (thus meets the requirement of LHS scaling tensor in DeepGEMM), this function will do nothing.
+
+    Arguments:
+        x: usually the LHS scaling tensor in GEMM.
+
+    Returns:
+        The LHS scaling tensor of TMA-aligned transposed format.
+    """
+    # NOTES: for the extreme performance, you may rewrite/fuse this function in CUDA
+    assert x.dim() in (2, 3)
+    remove_dim = False
+    if x.dim() == 2:
+        m, n = x.shape
+
+        aligned_m = get_tma_aligned_size(m, x.element_size())
+
+        if x.strides[0] == 1 and x.strides[1] == aligned_m:
+            return x
+
+        x, remove_dim = x.unsqueeze(0), True
+
+    b, m, n = x.shape
+    aligned_m = get_tma_aligned_size(m, x.element_size())
+
+    # The last kernel gives a column-major TMA aligned layout
+    if (
+        x.strides[0] == aligned_m * n
+        and x.strides[1] == 1
+        and x.strides[2] == aligned_m
+    ):
+        return x.squeeze(0) if remove_dim else x
+
+    # Normal layout requires transposing
+    aligned_x = paddle.transpose(
+        paddle.empty((b, n, aligned_m), dtype=x.dtype), perm=[0, 2, 1]
+    )
+    aligned_x[:, :m, :] = x
+    aligned_x = aligned_x[:, :m, :]
+    return aligned_x.squeeze(0) if remove_dim else aligned_x

python/paddle/incubate/fp8/deep_gemm/jit_kernels/wgrad_gemm.py

@@ -0,0 +1,241 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+from __future__ import annotations
+
+import os
+
+import paddle
+
+from ..jit import FP8WGradGemmRuntime, build
+from .gemm import get_best_configs
+from .runtime import (
+    GemmType,
+    make_2d_tma_a_desc,
+    make_2d_tma_b_desc,
+    make_2d_tma_d_desc,
+    make_2d_tma_scales_desc,
+)
+from .utils import (
+    ceil_div,
+    get_col_major_tma_aligned_tensor,
+    get_num_sms,
+    get_tma_aligned_size,
+)
+
+
+def wgrad_gemm_fp8_fp8_fp32_nt(
+    lhs: tuple[paddle.Tensor, paddle.Tensor],
+    rhs: tuple[paddle.Tensor, paddle.Tensor],
+    out: paddle.Tensor,
+    num_sms: int | None = None,
+):
+    """
+    Perform a weight gradient GEMM with FP8 inputs and FP32 output, with 1x128 LHS scaling and 1x128 RHS scaling.
+        Results will be accumulated into the output tensor.
+
+    Requirements:
+        LHS, RHS, and output tensors must be contiguous in dimension 1, i.e., strides[1] = 1.
+        The strides[0] of LHS and RHS must be a multiple of 16, and the strides[0] of output must be a multiple of 4.
+        RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling and RHS scaling tensor require a TMA-aligned transposed format.
+            If your input does not match the requirement, this function will do a transposing with a set of slow PaddlePaddle operations.
+
+    Arguments:
+        lhs: the first element is an FP8 tensor (typed `paddle.bfloat16`) of shape `[m, k]`,
+             the second element is an FP32 1x128 scaling tensor for LHS of shape `[m, ⌈k / 128⌉]`.
+        rhs: the first element is an FP8 tensor (typed `paddle.bfloat16`) of shape `[n, k]`,
+             the second element is an FP32 1x128 scaling tensor for RHS of shape `[n, ⌈k / 128⌉]`.
+        out: the FP32 output tensor of shape `[m, n]`, which will be accumulated.
+    """
+    lhs, lhs_scales = lhs
+    rhs, rhs_scales = rhs
+    m, k = lhs.shape
+    n, k_ = rhs.shape
+    m_, n_ = out.shape
+
+    # Type and shape checks
+    assert m == m_ and n == n_ and k == k_
+    assert n > 0 and m > 0
+    assert lhs_scales.shape == [m, ceil_div(k, 128)] or lhs_scales.shape == [
+        ceil_div(k, 128),
+        m,
+    ]
+    assert rhs_scales.shape == [n, ceil_div(k, 128)] or rhs_scales.shape == [
+        ceil_div(k, 128),
+        n,
+    ]
+    assert (
+        lhs.dtype == paddle.float8_e4m3fn and lhs_scales.dtype == paddle.float32
+    )
+    assert (
+        rhs.dtype == paddle.float8_e4m3fn and rhs_scales.dtype == paddle.float32
+    )
+    assert out.dtype == paddle.float32
+    assert lhs.strides[1] == 1 and out.strides[1] == 1 and rhs.strides[1] == 1
+
+    # LHS and RHS scales must be transposed for TMA load
+    # NOTES: `get_col_major_tma_aligned_tensor` may launch a kernel if not processed by previous kernels
+    def get_valid_scales(scales: paddle.Tensor, mn: int):
+        if scales.shape == [ceil_div(k, 128), mn]:
+            # For k-grouped GEMMs
+            scales = scales.transpose([1, 0])
+            assert get_tma_aligned_size(mn, 4) == scales.strides[1] == mn
+        else:
+            scales = get_col_major_tma_aligned_tensor(scales)
+        return scales
+
+    lhs_scales = get_valid_scales(lhs_scales, m)
+    rhs_scales = get_valid_scales(rhs_scales, n)
+
+    # Do nothing if `k` is zero
+    if k == 0:
+        return
+
+    # K must be aligned to 128
+    aligned_k = ceil_div(k, 128) * 128
+
+    # Auto-tuning with compilation
+    if num_sms is None:
+        num_sms = get_num_sms()
+    num_sms, block_m, block_n, num_stages, tma_multicast_config, smem_config = (
+        get_best_configs(
+            m, n, aligned_k, 1, num_sms, is_fp32_out=True, is_wgrad=True
+        )
+    )
+    if int(os.getenv("DG_JIT_KERNELS_DEBUG", 0)):
+        print(
+            f"Auto-tuned wgrad_gemm_fp8_fp8_fp32_nt as num_sms={num_sms}, block_m={block_m}, block_n={block_n}"
+        )
+    num_last_stages = ceil_div(k, 128) % num_stages
+    block_k = 128
+    num_tma_threads = 128
+    num_math_threads_per_group = 128
+
+    tensor_map_a = make_2d_tma_a_desc(
+        GemmType.Normal, lhs, m, k, lhs.strides[0], block_m, block_k, 1
+    )
+    tensor_map_b = make_2d_tma_b_desc(
+        GemmType.Normal, rhs, n, k, rhs.strides[0], block_n, block_k, 1
+    )
+    tensor_map_d = make_2d_tma_d_desc(
+        GemmType.Normal,
+        out,
+        m,
+        n,
+        out.strides[0],
+        block_m,
+        block_n,
+        1,
+        smem_config[1],
+    )
+    tensor_map_scales_a = make_2d_tma_scales_desc(
+        GemmType.Normal, lhs_scales, m, k, block_m, block_k, 1
+    )
+    tensor_map_scales_b = make_2d_tma_scales_desc(
+        GemmType.Normal, rhs_scales, n, k, block_n, block_k, 1
+    )
+
+    kwargs = {
+        # Templated arguments
+        "GEMM_TYPE": GemmType.Normal,
+        "NUM_TMA_THREADS": num_tma_threads,
+        "NUM_MATH_THREADS_PER_GROUP": num_math_threads_per_group,
+        "M": m,
+        "N": n,
+        "K": aligned_k,
+        "NUM_GROUPS": 1,
+        "BLOCK_M": block_m,
+        "BLOCK_N": block_n,
+        "BLOCK_K": block_k,
+        "NUM_STAGES": num_stages,
+        "NUM_LAST_STAGES": num_last_stages,
+        "NUM_TMA_MULTICAST": tma_multicast_config[0],
+        "IS_TMA_MULTICAST_ON_A": tma_multicast_config[1],
+        # Runtime arguments
+        "NUM_SMS": num_sms,
+        "SMEM_SIZE": smem_config[0],
+        "TENSOR_MAP_A": tensor_map_a,
+        "TENSOR_MAP_B": tensor_map_b,
+        "TENSOR_MAP_SCALES_A": tensor_map_scales_a,
+        "TENSOR_MAP_SCALES_B": tensor_map_scales_b,
+        "TENSOR_MAP_D": tensor_map_d,
+        "STREAM": paddle.device.current_stream().stream_base.cuda_stream,
+        "DEVICE_INDEX": out.place.gpu_device_id(),
+    }
+
+    # Generate, build and run the kernel
+    runtime = build("wgrad_gemm_fp8_fp8_fp32_nt", FP8WGradGemmRuntime, kwargs)
+    runtime(**kwargs)
+
+
+def k_grouped_wgrad_gemm_fp8_fp8_fp32_nt(
+    lhs: tuple[paddle.Tensor, paddle.Tensor],
+    rhs: tuple[paddle.Tensor, paddle.Tensor],
+    out: paddle.Tensor,
+    batch_sizes: list[int],
+    num_sms: int | None = None,
+):
+    """
+    Perform a k-grouped weight gradient GEMM with FP8 inputs and FP32 output, with 1x128 LHS scaling and 1x128 RHS scaling.
+        Results will be accumulated into the output tensor.
+
+    Requirements:
+        This function handles multiple batches with varying k-dimensions, processing each batch sequentially.
+        Each batch's LHS, RHS, and output tensors must be contiguous.
+        The RHS and RHS scaling factors are required to be transposed.
+        The LHS scaling and RHS scaling tensors require a TMA-aligned transposed format.
+
+    Arguments:
+        lhs: The first element is a flattened FP8 tensor (typed `paddle.bfloat16`) containing all batches of LHS data,
+                 and the flattened shape is `[sum(m * k for k in batch_sizes)]`, where m is the number of rows.
+             The second element is an FP32 scaling tensor for LHS with shape `[⌈k / 128⌉ for k in batch_sizes), m]`,
+                 representing the per-128-channel scaling factors.
+        rhs: The first element is a flattened FP8 tensor (typed `paddle.bfloat16`) containing all batches of RHS data,
+                 and the flattened shape is `[sum(n * k for k in batch_sizes)]`, where n is the number of rows.
+             The second element is an FP32 scaling tensor for RHS with shape `[⌈k / 128⌉ for k in batch_sizes), n]`,
+                 representing the per-128-channel scaling factors.
+        out: The FP32 output tensor of shape [num_batches, m, n], which will be accumulated.
+        batch_sizes: A list of integers specifying the k-dimension for each batch.
+    """
+    lhs, lhs_scales = paddle.view(lhs[0], [-1]), lhs[1]
+    rhs, rhs_scales = paddle.view(rhs[0], [-1]), rhs[1]
+    num_batches, m, n = out.shape
+
+    lhs_offset, rhs_offset, scales_offset = 0, 0, 0
+
+    for i in range(num_batches):
+        k = batch_sizes[i]
+        lhs_slice = paddle.view(lhs[lhs_offset : lhs_offset + m * k], (m, k))
+        rhs_slice = paddle.view(rhs[rhs_offset : rhs_offset + n * k], (n, k))
+        lhs_scales_slice = lhs_scales[
+            scales_offset : scales_offset + ceil_div(k, 128)
+        ]
+        rhs_scales_slice = rhs_scales[
+            scales_offset : scales_offset + ceil_div(k, 128)
+        ]
+        wgrad_gemm_fp8_fp8_fp32_nt(
+            (lhs_slice, lhs_scales_slice),
+            (rhs_slice, rhs_scales_slice),
+            out[i],
+            num_sms,
+        )
+
+        lhs_offset += m * k
+        rhs_offset += n * k
+        scales_offset += ceil_div(k, 128)

python/paddle/incubate/fp8/deep_gemm/utils.py

@@ -0,0 +1,82 @@
+# Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+#
+# 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.
+
+# The file has been adapted from DeepSeek DeepGEMM project
+# Copyright (c) 2025 DeepSeek
+# Licensed under the MIT License - https://github.com/deepseek-ai/DeepGEMM/blob/main/LICENSE
+
+import os
+import sys
+
+import paddle
+
+
+class empty_suppress:
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *_):
+        pass
+
+
+class suppress_stdout_stderr:
+    def __enter__(self):
+        self.outnull_file = open(os.devnull, "w")
+        self.errnull_file = open(os.devnull, "w")
+
+        self.old_stdout_fileno_undup = sys.stdout.fileno()
+        self.old_stderr_fileno_undup = sys.stderr.fileno()
+
+        self.old_stdout_fileno = os.dup(sys.stdout.fileno())
+        self.old_stderr_fileno = os.dup(sys.stderr.fileno())
+
+        self.old_stdout = sys.stdout
+        self.old_stderr = sys.stderr
+
+        os.dup2(self.outnull_file.fileno(), self.old_stdout_fileno_undup)
+        os.dup2(self.errnull_file.fileno(), self.old_stderr_fileno_undup)
+
+        sys.stdout = self.outnull_file
+        sys.stderr = self.errnull_file
+        return self
+
+    def __exit__(self, *_):
+        sys.stdout = self.old_stdout
+        sys.stderr = self.old_stderr
+
+        os.dup2(self.old_stdout_fileno, self.old_stdout_fileno_undup)
+        os.dup2(self.old_stderr_fileno, self.old_stderr_fileno_undup)
+
+        os.close(self.old_stdout_fileno)
+        os.close(self.old_stderr_fileno)
+
+        self.outnull_file.close()
+        self.errnull_file.close()
+
+
+def calc_diff(x, y):
+    x, y = x.astype(paddle.float64), y.astype(paddle.float64)
+    denominator = (x * x + y * y).sum()
+    sim = 2 * (x * y).sum() / denominator
+    return 1 - sim
+
+
+def count_bytes(tensors):
+    total = 0
+    for t in tensors:
+        if isinstance(t, tuple):
+            total += count_bytes(t)
+        else:
+            total += t.numel() * t.element_size()
+    return total