DeepGEMM/third-party/tilelang_ops/swiglu_apply_weight_to_fp8.py

import deep_gemm
import tilelang
import torch
from math import gcd
from tilelang import language as T

from .utils import get_sf_and_inv, get_sf_shape


@tilelang.jit
def _swiglu_apply_weight_to_fp8_tl(
    half_hidden: int,
    num_per_channels: int,
    use_col_major_scales: bool,
    round_scale: bool,
    ue8m0_scale: bool,
    num_ctas: int,
    has_topk_weights: bool,
    has_avail_tokens: bool,
    has_clamp_value: bool,
    output_bf16: bool,
    fast_math: bool,
) -> None:
    in_dtype = T.bfloat16
    w_dtype = T.float32
    out_dtype = T.float8_e4m3fn
    out_sf_dtype = T.uint8 if ue8m0_scale else T.float32
    num_dtype = T.int32
    num_tokens = T.dynamic("num_tokens")
    assert half_hidden % 16 == 0
    if not output_bf16:
        assert half_hidden % num_per_channels == 0

    num_block_h = max(half_hidden // gcd(half_hidden, 16 * 1024), 1)
    assert num_block_h <= num_ctas, "not supported hidden size"
    blk_h = half_hidden // num_block_h

    layout_h = blk_h // 16

    blk_n = 1024 // layout_h

    def local_layout(i, j):  # noqa: ANN001
        thread_id = i * layout_h + j // 16
        local_id = j % 16
        return thread_id, local_id

    def local_layout_3d(i, j, k):  # noqa: ANN001
        return local_layout(i, j * num_per_channels + k)

    @T.macro
    def main(
        bi: int,
        bh: int,
        num_ctas: int,
        x: T.Tensor[(num_tokens, half_hidden * 2), in_dtype],  # type: ignore
        topk_weights: T.Tensor[num_tokens, w_dtype],  # type: ignore
        avail_tokens: T.Tensor[1, num_dtype],  # type: ignore
        out: T.Tensor[(num_tokens, half_hidden), out_dtype],  # type: ignore
        out_sf: T.Tensor[get_sf_shape(num_tokens, half_hidden, num_per_channels, ue8m0_scale, use_col_major_scales), out_sf_dtype],  # type: ignore
        out_bf16: T.Tensor[(num_tokens, half_hidden), T.bfloat16],  # type: ignore
        clamp_value: T.float32,
    ):
        gate_frag = T.alloc_fragment((blk_n, blk_h), T.float32)
        up_frag = T.alloc_fragment((blk_n, blk_h), T.float32)
        y_frag = T.alloc_fragment((blk_n, blk_h // num_per_channels, num_per_channels), T.float32)
        y_f8_frag = T.alloc_fragment((blk_n, blk_h), out_dtype)
        T.annotate_layout(
            {
                gate_frag: T.Fragment(gate_frag.shape, forward_fn=local_layout),
                up_frag: T.Fragment(up_frag.shape, forward_fn=local_layout),
                y_frag: T.Fragment(y_frag.shape, forward_fn=local_layout_3d),
                y_f8_frag: T.Fragment(y_f8_frag.shape, forward_fn=local_layout),
            }
        )

        T.assume(0 <= bh * blk_h + blk_h <= half_hidden)

        for i, j in T.Parallel(blk_n, blk_h):
            gate_frag[i, j] = x[bi + i * num_ctas, bh * blk_h + j]
        for i, j in T.Parallel(blk_n, blk_h):
            up_frag[i, j] = x[bi + i * num_ctas, half_hidden + bh * blk_h + j]

        topk_weight = T.alloc_fragment((blk_n,), T.float32)
        for i in T.Parallel(blk_n):
            topk_weight[i] = topk_weights[bi + i * num_ctas] if has_topk_weights else 1.0

        zero = T.alloc_var(T.float32, 0.0)
        for i, j in T.Parallel(blk_n, blk_h):
            if has_clamp_value:
                up_frag[i, j] = T.min(clamp_value, T.max(-clamp_value, up_frag[i, j]))
                gate_frag[i, j] = T.min(clamp_value, gate_frag[i, j])
            y_frag[i, j // num_per_channels, j % num_per_channels] = (
                gate_frag[i, j] / (1 + T.exp(-gate_frag[i, j])) * up_frag[i, j] * topk_weight[i] + zero
            )  # HACK : + 0 for vectorize

        y_max_frag = T.alloc_fragment((blk_n, blk_h // num_per_channels), T.float32)
        sf_inv_frag = T.alloc_fragment((blk_n, blk_h // num_per_channels), T.float32)
        T.reduce_absmax(T.reshape(y_frag, (blk_n, blk_h // num_per_channels, num_per_channels)), y_max_frag)
        for i, j in T.Parallel(blk_n, blk_h // num_per_channels):
            clamped_amax = T.max(y_max_frag[i, j], 1e-4)
            sf, sf_inv = get_sf_and_inv(clamped_amax, round_scale, ue8m0_scale)
            i_index = bi + i * num_ctas
            j_index = blk_h // num_per_channels * bh + j
            # Store SF
            if ue8m0_scale:
                out_sf[j_index // 4, i_index * 4 + j_index % 4] = sf
            elif use_col_major_scales:
                out_sf[j_index, i_index] = sf
            else:
                out_sf[i_index, j_index] = sf
            sf_inv_frag[i, j] = sf_inv

        for i, j in T.Parallel(blk_n, blk_h):
            y_f8_frag[i, j] = y_frag[i, j // num_per_channels, j % num_per_channels] * sf_inv_frag[i, j // num_per_channels]

        for i, j in T.Parallel(blk_n, blk_h):
            out[bi + i * num_ctas, blk_h * bh + j] = y_f8_frag[i, j]

        if output_bf16:
            for i, j in T.Parallel(blk_n, blk_h):
                out_bf16[bi + i * num_ctas, blk_h * bh + j] = y_frag[i, j // num_per_channels, j % num_per_channels]

    @T.prim_func
    def _swiglu_apply_weight_to_fp8(
        x: T.Tensor[(num_tokens, half_hidden * 2), in_dtype],  # type: ignore
        topk_weights: T.Tensor[num_tokens, w_dtype],  # type: ignore
        avail_tokens: T.Tensor[1, num_dtype],  # type: ignore
        out: T.Tensor[(num_tokens, half_hidden), out_dtype],  # type: ignore
        out_sf: T.Tensor[get_sf_shape(num_tokens, half_hidden, num_per_channels, ue8m0_scale, use_col_major_scales), out_sf_dtype],  # type: ignore
        out_bf16: T.Tensor[(num_tokens, half_hidden), T.bfloat16],  # type: ignore
        clamp_value: T.float32,
    ):
        # we actually don't use this
        _ = num_tokens
        # simplest schedule: one token one block, but pipelined as persistent
        with T.Kernel(num_ctas, threads=1024) as cta_id:
            avail_tokens_l = avail_tokens[0] if has_avail_tokens else num_tokens
            T.pdl_sync()  # avail_tokens must be const.
            T.assume(0 <= avail_tokens_l <= num_tokens)
            thread_idx = T.get_thread_binding()
            new_num_ctas = num_ctas // num_block_h
            if cta_id >= new_num_ctas * num_block_h:
                T.thread_return()
            for bi in T.serial(cta_id // num_block_h, avail_tokens_l - thread_idx // layout_h * new_num_ctas, new_num_ctas * blk_n):
                main(bi, cta_id % num_block_h, new_num_ctas, x, topk_weights, avail_tokens, out, out_sf, out_bf16, clamp_value)

    return _swiglu_apply_weight_to_fp8


def swiglu_apply_weight_to_fp8(
    x: torch.Tensor,
    topk_weights: torch.Tensor | None,
    avail_tokens: torch.Tensor | None,
    num_per_channels: int,
    use_col_major_scales: bool,
    round_scale: bool,
    ue8m0_scale: bool,
    clamp_value: float | None = None,
    fmt: str = "e4m3",
    num_sms: int | None = None,
    output_bf16: bool = False,
    fast_math: bool = True,
) -> tuple[torch.Tensor, torch.Tensor] | tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
    assert fmt == "e4m3"
    if num_sms is None:
        num_sms = deep_gemm.get_num_sms()

    num_tokens, hidden_size = x.shape
    assert hidden_size % (2 * num_per_channels) == 0

    y = torch.empty(
        (num_tokens, hidden_size // 2),
        device=x.device,
        dtype=torch.float8_e4m3fn,
    )
    y_sf = torch.empty(
        get_sf_shape(num_tokens, hidden_size // 2, num_per_channels, ue8m0_scale, use_col_major_scales),
        device=x.device,
        dtype=(torch.uint8 if ue8m0_scale else torch.float32),
    )

    y_bf16 = torch.empty((num_tokens, hidden_size // 2), device=x.device, dtype=torch.bfloat16) if output_bf16 else None

    if num_tokens > 0:
        _swiglu_apply_weight_to_fp8_tl.pass_configs = {
            tilelang.PassConfigKey.TL_DISABLE_WARP_SPECIALIZED: True,
            tilelang.PassConfigKey.TL_DISABLE_TMA_LOWER: True,
            tilelang.PassConfigKey.TL_ENABLE_FAST_MATH: fast_math,
        }
        kernel = _swiglu_apply_weight_to_fp8_tl(
            hidden_size // 2,
            num_per_channels,
            use_col_major_scales,
            round_scale,
            ue8m0_scale,
            num_sms,
            topk_weights is not None,
            avail_tokens is not None,
            clamp_value is not None,
            output_bf16,
            fast_math
        )
        kernel(x, topk_weights, avail_tokens.view(1) if avail_tokens is not None else None, y, y_sf, y_bf16, clamp_value or 0.0)

    if ue8m0_scale:
        if num_tokens == 0:
            y_sf.as_strided_(y_sf.size(), (0, 1))
        y_sf = y_sf.view(dtype=torch.int32)
    if output_bf16:
        return y, y_sf.T[:num_tokens], y_bf16
    else:
        return y, y_sf.T[:num_tokens]