mlc-ai--mlc-llm/python/mlc_llm/compiler_pass/attach_softmax_with_temperature.py

"""A compiler pass that attaches two-stage softmax with temperature."""

from typing import Any, Dict, Optional  # noqa: UP035

import tvm
from tvm import relax, tirx
from tvm.ir.module import IRModule
from tvm.relax.expr_functor import PyExprMutator, mutator
from tvm.script import tirx as T

from ..support.max_thread_check import get_max_num_threads_per_block


@tvm.transform.module_pass(opt_level=0, name="AttachSoftmaxWithTemperature")
class AttachSoftmaxWithTemperature:
    """Rewrites one-shot softmax into two-stage softmax."""

    def __init__(
        self,
        target: tvm.target.Target,
        metadata: Optional[Dict[str, Any]] = None,  # noqa: UP006
    ) -> None:
        self.target = target
        self.metadata = metadata

    def transform_module(self, mod: IRModule, _ctx: tvm.transform.PassContext) -> IRModule:
        """IRModule-level transformation"""
        return _Rewriter(mod, self.target, self.metadata).transform()


@mutator
class _Rewriter(PyExprMutator):
    def __init__(
        self,
        mod: IRModule,
        target: tvm.target.Target,
        metadata: Optional[Dict[str, Any]] = None,  # noqa: UP006
    ) -> None:
        super().__init__(mod)
        self.mod = mod
        self.target = target
        self.metadata = metadata
        self.chunk_size = 4096
        self.active_vocab_size = self.metadata.get("active_vocab_size") if self.metadata else None

    def transform(self) -> IRModule:
        """Entry point"""
        batch_size = tirx.Var("batch_size", "int64")
        vocab_size = tirx.Var("vocab_size", "int64")
        dtype = "float32"
        logits = relax.Var("logits", relax.TensorType([batch_size, 1, vocab_size], dtype))
        temperature = relax.Var("temperature", relax.TensorType([batch_size], dtype))
        with self.builder_.function("softmax_with_temperature", params=[logits, temperature]):
            with self.builder_.dataflow():
                output_struct_info = logits.ty
                new_shape = relax.ShapeExpr([batch_size, vocab_size])
                logits = relax.call_pure_packed(
                    "vm.builtin.reshape",
                    logits,
                    new_shape,
                    ty_args=relax.TensorType(new_shape, dtype),
                )
                f_chunk_lse, f_softmax_with_lse = _get_lse_and_softmax_func(
                    self.target, self.chunk_size, self.active_vocab_size
                )
                chunked_result_struct_info = relax.TensorType(
                    (batch_size, (vocab_size + self.chunk_size - 1) // self.chunk_size),
                    "float32",
                )
                chunked_results = self.builder_.emit(
                    relax.call_tir(
                        self.builder_.add_func(f_chunk_lse, "chunk_lse"),
                        args=[logits, temperature],
                        out_ty=[
                            chunked_result_struct_info,
                            chunked_result_struct_info,
                        ],
                    )
                )
                chunked_sum = chunked_results[0]
                chunked_max = chunked_results[1]
                softmax = self.builder_.emit(
                    relax.call_tir(
                        self.builder_.add_func(f_softmax_with_lse, "softmax_with_chunked_sum"),
                        args=[logits, temperature, chunked_sum, chunked_max],
                        out_ty=logits.ty,
                    )
                )
                softmax = self.builder_.emit_output(
                    relax.call_pure_packed(
                        "vm.builtin.reshape",
                        softmax,
                        output_struct_info.shape,
                        ty_args=output_struct_info,
                    )
                )
            self.builder_.emit_func_output(softmax)
        return self.builder_.get()


def _get_lse_and_softmax_func(target: tvm.target.Target, chunk_size: int, active_vocab_size: int):
    # NOTE: A quick note on the softmax implementation.
    # We once tried to multiply every element by log2e which can be computed
    # potentially more efficiently on hardware.
    # However, when the input values are large, multiplying by the factor of log2e
    # causes numerical issue in float32 dtype.
    # This leads to the softmax output not summing up to 1.
    # For numerical stability, we removed the log2e factor and switched back
    # to the standard log/exp computation.

    # The kernels below handle both the cases of temperature=0 and temperature != 0.
    # - When temperature is not 0, the first kernel computes the log-sum-exp of
    # chunks (subtracted by the max value in chunk), and the max values of chunks.
    # The second kernel merges the log-sum-exp with the maximum values.
    # - When temperature is 0, the first kernel computes the max value and the counts
    # of the max value. The second kernel merges the max and counts, and set the
    # softmax of the maximum values to "max_value / max_count".

    @T.prim_func(s_tir=True)
    def chunk_lse(
        var_A: T.handle,
        var_temperature: T.handle,
        var_chunked_sum: T.handle,
        var_chunked_max: T.handle,
    ):
        T.func_attr({"tirx.noalias": True})
        batch_size = T.int64()
        vocab_size = T.int64()
        num_chunks = T.int64()
        A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32")
        temperature = T.match_buffer(var_temperature, (batch_size,), dtype="float32")
        chunked_sum = T.match_buffer(var_chunked_sum, (batch_size, num_chunks), dtype="float32")
        chunked_max = T.match_buffer(var_chunked_max, (batch_size, num_chunks), dtype="float32")
        A_pad = T.sblock_alloc_buffer(
            (batch_size, num_chunks, T.int64(chunk_size)), dtype="float32"
        )
        temp_max = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32")
        temp_sum = T.sblock_alloc_buffer((batch_size, num_chunks), dtype="float32")

        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
            with T.sblock("pad"):
                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
                A_pad[v0, v1, v2] = T.Select(
                    v1 * T.int64(chunk_size) + v2
                    < (active_vocab_size if active_vocab_size is not None else vocab_size),
                    T.if_then_else(
                        temperature[v0] > T.float32(1e-5),
                        A[v0, v1 * T.int64(chunk_size) + v2] / temperature[v0],
                        A[v0, v1 * T.int64(chunk_size) + v2],
                    ),
                    T.min_value("float32"),
                )
        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
            with T.sblock("max"):
                v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2])
                with T.init():
                    temp_max[v0, v1] = T.min_value("float32")
                temp_max[v0, v1] = T.max(temp_max[v0, v1], A_pad[v0, v1, v2])
        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
            with T.sblock("sum_exp"):
                v0, v1, v2 = T.axis.remap("SSR", [l0, l1, l2])
                with T.init():
                    temp_sum[v0, v1] = T.float32(0)
                temp_sum[v0, v1] += T.if_then_else(
                    v1 * T.int64(chunk_size) + v2
                    < (active_vocab_size if active_vocab_size is not None else vocab_size),
                    T.Select(
                        temperature[v0] > T.float32(1e-5),
                        T.exp(A_pad[v0, v1, v2] - temp_max[v0, v1]),
                        T.cast(A_pad[v0, v1, v2] == temp_max[v0, v1], "float32"),
                    ),
                    T.float32(0),
                )
        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(1)):
            with T.sblock("log"):
                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
                chunked_sum[v0, v1] = T.Select(
                    temperature[v0] > T.float32(1e-5),
                    T.log(temp_sum[v0, v1]),
                    temp_sum[v0, v1],
                )
                chunked_max[v0, v1] = temp_max[v0, v1]

    @T.prim_func(s_tir=True)
    def softmax_with_chunked_sum(
        var_A: T.handle,
        var_temperature: T.handle,
        var_chunked_sum: T.handle,
        var_chunked_max: T.handle,
        var_softmax: T.handle,
    ):
        T.func_attr({"tirx.noalias": True, "tirx.is_scheduled": 1})
        batch_size = T.int64()
        vocab_size = T.int64()
        num_chunks = T.int64()
        A = T.match_buffer(var_A, (batch_size, vocab_size), dtype="float32")
        temperature = T.match_buffer(var_temperature, (batch_size,), dtype="float32")
        chunked_sum = T.match_buffer(var_chunked_sum, (batch_size, num_chunks), dtype="float32")
        chunked_max = T.match_buffer(var_chunked_max, (batch_size, num_chunks), dtype="float32")
        softmax = T.match_buffer(var_softmax, (batch_size, vocab_size), dtype="float32")
        temp_max = T.sblock_alloc_buffer((batch_size,), dtype="float32")
        temp_sum = T.sblock_alloc_buffer((batch_size,), dtype="float32")
        for l0, l1 in T.grid(batch_size, num_chunks):
            with T.sblock("max"):
                v0, v1 = T.axis.remap("SR", [l0, l1])
                with T.init():
                    temp_max[v0] = T.min_value("float32")
                temp_max[v0] = T.max(temp_max[v0], chunked_max[v0, v1])
        for l0, l1 in T.grid(batch_size, num_chunks):
            with T.sblock("sum_exp"):
                v0, v1 = T.axis.remap("SR", [l0, l1])
                with T.init():
                    temp_sum[v0] = T.float32(0)
                temp_sum[v0] += T.Select(
                    temperature[v0] > T.float32(1e-5),
                    T.exp(chunked_sum[v0, v1] + chunked_max[v0, v1] - temp_max[v0]),
                    T.cast(chunked_max[v0, v1] == temp_max[v0], "float32") * chunked_sum[v0, v1],
                )
        for l0, l1, l2 in T.grid(batch_size, num_chunks, T.int64(chunk_size)):
            with T.sblock("log_pad"):
                v0, v1, v2 = T.axis.remap("SSS", [l0, l1, l2])
                if v1 * T.int64(chunk_size) + v2 < vocab_size:
                    softmax[v0, v1 * T.int64(chunk_size) + v2] = T.Select(
                        v1 * T.int64(chunk_size) + v2
                        < (active_vocab_size if active_vocab_size is not None else vocab_size),
                        T.if_then_else(
                            temperature[v0] > T.float32(1e-5),
                            T.exp(
                                A[v0, v1 * T.int64(chunk_size) + v2] / temperature[v0]
                                - (T.log(temp_sum[v0]) + temp_max[v0])
                            ),
                            T.cast(
                                A[v0, v1 * T.int64(chunk_size) + v2] == temp_max[v0],
                                "float32",
                            )
                            / temp_sum[v0],
                        ),
                        T.float32(0),
                    )

    sch = tvm.s_tir.Schedule(IRModule({"softmax_with_chunked_sum": softmax_with_chunked_sum}))

    def apply_gpu_schedule(target, sch):
        max_threads = get_max_num_threads_per_block(target)
        TX = 32
        TY = max_threads // TX
        unroll_depth = 64

        sch.work_on("softmax_with_chunked_sum")
        l0, l1, l2 = sch.get_loops("log_pad")
        bx = sch.fuse(l0, l1)
        sch.bind(bx, "blockIdx.x")
        unroll, ty, tx = sch.split(l2, [None, TY, TX])
        sch.bind(ty, "threadIdx.y")
        sch.bind(tx, "threadIdx.x")
        sch.annotate(unroll, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
        sch.annotate(unroll, ann_key="pragma_unroll_explicit", ann_val=1)

        for block_name in ["sum_exp", "max"]:
            block = sch.get_sblock(block_name)
            sch.set_scope(block, buffer_index=0, storage_scope="shared")
            sch.compute_at(block, bx)
            r_loop = sch.get_loops(block)[-1]
            r_loop, tx = sch.split(r_loop, [None, TX])
            sch.reorder(tx, r_loop)
            sch.bind(tx, "threadIdx.x")
            sch.annotate(r_loop, ann_key="pragma_auto_unroll_max_step", ann_val=unroll_depth)
            sch.annotate(r_loop, ann_key="pragma_unroll_explicit", ann_val=1)

        return chunk_lse, sch.mod["softmax_with_chunked_sum"]

    if target.kind.name == "llvm":
        return chunk_lse, sch.mod["softmax_with_chunked_sum"]
    return apply_gpu_schedule(target, sch)