chore: import upstream snapshot with attribution

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)