generate_kernels.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import glob
import itertools
import os
import subprocess
import sys

import jinja2

ARCHS = []
SUPPORT_FP8 = False
SUPPORT_SM75 = False
SUPPORT_SM80 = False
for arch in sys.argv[1].split(","):
    arch = arch[: arch.index(".") + 2].replace(".", "")
    arch = int(arch)
    # only SM89 and SM120 fully support
    # mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32.
    # SM90 and SM100 can use this PTX, but it’s simulated
    # with FP16 MMA, so it cannot achieve any acceleration.
    if arch in [89, 120]:
        SUPPORT_FP8 = True
    if arch >= 80:
        SUPPORT_SM80 = True
    if arch == 75:
        SUPPORT_SM75 = True

FILE_HEAD_COMMENT = """
// auto generated by generate_kernels.py
// clang-format off
""".lstrip()

FILE_HEAD = (
    FILE_HEAD_COMMENT
    + """
#include "kernel.h"
#include "marlin_template.h"

namespace MARLIN_NAMESPACE_NAME {
"""
)

TEMPLATE = (
    "template __global__ void Marlin<"
    "{{a_type_id}}, "
    "{{b_type_id}}, "
    "{{c_type_id}}, "
    "{{s_type_id}}, "
    "{{threads}}, "
    "{{thread_m_blocks}}, "
    "{{thread_n_blocks}}, "
    "{{thread_k_blocks}}, "
    "{{m_block_size_8}}, "
    "{{stages}}, "
    "{{group_blocks}}, "
    "{{is_zp_float}}>"
    "( MARLIN_KERNEL_PARAMS );"
)

THREAD_CONFIGS = [(128, 128, 256), (64, 256, 256), (64, 128, 128), (128, 64, 128)]

THREAD_M_BLOCKS = [0.5, 1, 2, 3, 4]

QUANT_CONFIGS = [
    # AWQ-INT4
    {
        "b_type": "kU4",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": THREAD_M_BLOCKS,
        "group_blocks": [-1, 2, 4, 8],
    },
    # GPTQ-INT4
    {
        "b_type": "kU4B8",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": THREAD_M_BLOCKS,
        "group_blocks": [-1, 0, 2, 4, 8],
    },
    # GPTQ-INT8
    {
        "b_type": "kU8B128",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": THREAD_M_BLOCKS,
        "group_blocks": [-1, 0, 2, 4, 8],
    },
    # FP8
    {
        "b_type": "kFE4M3fn",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": THREAD_M_BLOCKS,
        "group_blocks": [-1, 8],
    },
    # NVFP4
    {
        "b_type": "kFE2M1f",
        "s_type": "kFE4M3fn",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": THREAD_M_BLOCKS,
        "group_blocks": [1],
    },
    # MXFP4
    {
        "a_type": ["kBFloat16"],
        "b_type": "kFE2M1f",
        "s_type": "kFE8M0fnu",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": THREAD_M_BLOCKS,
        "group_blocks": [2],
    },
    # MXFP8
    {
        "a_type": ["kBFloat16"],
        "b_type": "kFE4M3fn",
        "s_type": "kFE8M0fnu",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": THREAD_M_BLOCKS,
        "group_blocks": [2],
    },
    # AWQ-INT4 with INT8 activation
    {
        "a_type": ["kS8"],
        "b_type": "kU4",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": [1, 2, 3, 4],
        "group_blocks": [-1, 2, 4, 8],
    },
    # GPTQ-INT4 with INT8 activation
    {
        "a_type": ["kS8"],
        "b_type": "kU4B8",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": [1, 2, 3, 4],
        "group_blocks": [-1, 2, 4, 8],
    },
    # GPTQ-INT4 with FP8 activation
    {
        "a_type": ["kFE4M3fn"],
        "b_type": "kU4B8",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": [1, 2, 3, 4],
        "group_blocks": [-1, 2, 4, 8],
    },
    # AWQ-INT4 with FP8 activation
    {
        "a_type": ["kFE4M3fn"],
        "b_type": "kU4",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": [1, 2, 3, 4],
        "group_blocks": [-1, 2, 4, 8],
    },
    # MXFP4 with FP8 activation
    {
        "a_type": ["kFE4M3fn"],
        "b_type": "kFE2M1f",
        "c_type": ["kBFloat16"],
        "s_type": "kFE8M0fnu",
        "thread_configs": THREAD_CONFIGS,
        "thread_m_blocks": [1, 2, 3, 4],
        "group_blocks": [2],
    },
]


def remove_old_kernels():
    for filename in glob.glob(os.path.dirname(__file__) + "/*kernel_*.cu"):
        subprocess.call(["rm", "-f", filename])

    filename = os.path.dirname(__file__) + "/kernel_selector.h"
    subprocess.call(["rm", "-f", filename])


def generate_new_kernels():
    result_dict = {}
    sm_75_result_dict = {}

    for quant_config in QUANT_CONFIGS:
        c_types = quant_config.get("c_type", ["kFloat16", "kBFloat16"])
        a_types = quant_config.get("a_type", ["kFloat16", "kBFloat16"])
        b_type = quant_config["b_type"]
        is_zp_float = quant_config.get("is_zp_float", False)
        all_group_blocks = quant_config["group_blocks"]
        all_m_blocks = quant_config["thread_m_blocks"]
        all_thread_configs = quant_config["thread_configs"]

        for a_type, c_type in itertools.product(a_types, c_types):
            if not SUPPORT_FP8 and a_type == "kFE4M3fn":
                continue
            if "16" in a_type and "16" in c_type and a_type != c_type:
                continue
            s_type = quant_config.get("s_type", c_type)
            if (a_type, b_type, c_type) not in result_dict:
                result_dict[(a_type, b_type, c_type)] = []
                if a_type in ["kFloat16", "kS8"] and c_type == "kFloat16":
                    sm_75_result_dict[(a_type, b_type, c_type)] = []

            for group_blocks, m_blocks, thread_configs in itertools.product(
                all_group_blocks, all_m_blocks, all_thread_configs
            ):
                thread_k, thread_n, threads = thread_configs

                if threads == 256:
                    # for small batch (m_blocks == 1),
                    #     we only need (128, 128, 256)
                    # for large batch (m_blocks > 1),
                    #     we only need (64, 256, 256)
                    if m_blocks <= 1 and (thread_k, thread_n) != (128, 128):
                        continue
                    if m_blocks > 1 and (thread_k, thread_n) != (64, 256):
                        continue

                config = {
                    "threads": threads,
                    "s_type": s_type,
                    "thread_m_blocks": max(m_blocks, 1),
                    "thread_k_blocks": thread_k // 16,
                    "thread_n_blocks": thread_n // 16,
                    "m_block_size_8": "true" if m_blocks == 0.5 else "false",
                    "stages": 4,
                    "group_blocks": group_blocks,
                    "is_zp_float": "true" if is_zp_float else "false",
                }

                if SUPPORT_SM80:
                    result_dict[(a_type, b_type, c_type)].append(config)
                if (a_type, b_type, c_type) in sm_75_result_dict and SUPPORT_SM75:
                    config_sm75 = config.copy()
                    config_sm75["stages"] = 2
                    sm_75_result_dict[(a_type, b_type, c_type)].append(config_sm75)

    kernel_selector_str = FILE_HEAD_COMMENT

    for result_dict_tmp in [result_dict, sm_75_result_dict]:
        for (a_type, b_type, c_type), config_list in result_dict_tmp.items():
            all_template_str_list = []
            if not config_list:
                continue
            for config in config_list:
                s_type = config["s_type"]
                template_str = jinja2.Template(TEMPLATE).render(
                    a_type_id=f"vllm::{a_type}.id()",
                    b_type_id=f"vllm::{b_type}.id()",
                    c_type_id=f"vllm::{c_type}.id()",
                    s_type_id=f"vllm::{s_type}.id()",
                    **config,
                )
                all_template_str_list.append(template_str)

                conditions = [
                    f"a_type == vllm::{a_type}",
                    f"b_type == vllm::{b_type}",
                    f"c_type == vllm::{c_type}",
                    f"s_type == vllm::{s_type}",
                    f"threads == {config['threads']}",
                    f"thread_m_blocks == {config['thread_m_blocks']}",
                    f"thread_n_blocks == {config['thread_n_blocks']}",
                    f"thread_k_blocks == {config['thread_k_blocks']}",
                    f"m_block_size_8 == {config['m_block_size_8']}",
                    f"stages == {config['stages']}",
                    f"group_blocks == {config['group_blocks']}",
                    f"is_zp_float == {config['is_zp_float']}",
                ]
                conditions = " && ".join(conditions)

                if kernel_selector_str == FILE_HEAD_COMMENT:
                    kernel_selector_str += f"if ({conditions})\n  kernel = "
                else:
                    kernel_selector_str += f"else if ({conditions})\n  kernel = "

                kernel_template2 = (
                    "Marlin<{{a_type_id}}, {{b_type_id}}, {{c_type_id}}, "
                    "{{s_type_id}}, {{threads}}, {{thread_m_blocks}}, "
                    "{{thread_n_blocks}}, {{thread_k_blocks}}, "
                    "{{m_block_size_8}}, {{stages}}, {{group_blocks}}, "
                    "{{is_zp_float}}>;"
                )

                kernel_selector_str += (
                    jinja2.Template(kernel_template2).render(
                        a_type_id=f"vllm::{a_type}.id()",
                        b_type_id=f"vllm::{b_type}.id()",
                        c_type_id=f"vllm::{c_type}.id()",
                        s_type_id=f"vllm::{s_type}.id()",
                        **config,
                    )
                    + "\n"
                )

            file_content = FILE_HEAD + "\n\n"
            file_content += "\n\n".join(all_template_str_list) + "\n\n}\n"
            if a_type == "kFE4M3fn":
                filename = f"sm89_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
            elif result_dict_tmp is sm_75_result_dict:
                filename = f"sm75_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"
            else:
                filename = f"sm80_kernel_{a_type[1:]}_{b_type[1:]}_{c_type[1:]}.cu"

            filename = filename.lower()

            with open(os.path.join(os.path.dirname(__file__), filename), "w") as f:
                f.write(file_content)

    if not SUPPORT_FP8 and kernel_selector_str != FILE_HEAD_COMMENT:
        kernel_selector_str += (
            "else if (a_type == vllm::kFE4M3fn)\n"
            "  TORCH_CHECK(false, "
            '"marlin kernel with fp8 activation is not built.");'
        )

    with open(os.path.join(os.path.dirname(__file__), "kernel_selector.h"), "w") as f:
        f.write(kernel_selector_str)


if __name__ == "__main__":
    remove_old_kernels()
    generate_new_kernels()