TensorRT-LLMs/tests/integration/defs/examples/test_openai.py

# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
"""Module test_openai test openai examples."""
import os
import subprocess  # fmt: off

import pytest
from defs.common import find_tensorrt, venv_check_call
from defs.trt_test_alternative import call, check_call, make_clean_dirs


@pytest.fixture(scope="module")
def openai_triton_example_root(llm_root):
    "Get openai-triton example root"
    example_root = os.path.join(llm_root, "examples", "openai_triton",
                                "manual_plugin")

    return example_root


@pytest.fixture(scope="module")
def openai_triton_plugingen_example_root(llm_root):
    "Get openai-triton PluginGen example root"
    example_root = os.path.join(llm_root, "examples", "openai_triton",
                                "plugin_autogen")

    return example_root


@pytest.fixture(scope="module")
def llm_openai_triton_model_root(llm_venv):
    "prepare openai-triton model & return model root"
    workspace = llm_venv.get_working_directory()
    model_root = os.path.join(workspace, "triton")
    commit = "d4644d6cb3ae674e1f15932cac1f28104795744f"

    call(f"git clone https://github.com/openai/triton.git {model_root}",
         shell=True)
    call(f"cd {model_root} && git checkout {commit}", shell=True)
    llm_venv.run_cmd(["-m", "pip", "install", "cmake"])
    llm_venv.run_cmd([
        "-m", "pip", "install",
        os.path.abspath(os.path.join(model_root, "python"))
    ])

    yield model_root

    llm_venv.run_cmd(["-m", "pip", "uninstall", "-y", "triton"])


def test_llm_openai_triton_1gpu(openai_triton_example_root,
                                llm_openai_triton_model_root, llm_venv,
                                engine_dir, trt_config, is_trt_environment):
    aot_path = os.path.join(openai_triton_example_root, "aot")
    aot_fp16_path = os.path.join(aot_path, "fp16")
    aot_fp32_path = os.path.join(aot_path, "fp32")
    call(f"mkdir -p {aot_fp16_path}", shell=True)
    call(f"mkdir -p {aot_fp32_path}", shell=True)

    num_stages = "2"

    # yapf: disable
    # Kernel for data type=float16, BLOCK_M=128, BLOCK_DMODEL=64, BLOCK_N=128
    compile_cmd = [
        f"{llm_openai_triton_model_root}/python/triton/tools/compile.py",
        f"{openai_triton_example_root}/fmha_triton.py",
        "-n", "fused_attention_kernel",
        "-o", f"{aot_fp16_path}/fmha_kernel_d64_fp16",
        "--out-name", "fmha_d64_fp16", "-w", "4", "-ns", num_stages,
        "-s", "*fp16:16, *fp32:16, *fp32:16, *fp16:16, *fp16:16, *fp16:16, fp32, i32, i32, i32, 128, 64, 128",
        "-g", "(seq_len + 127) / 128, batch_size * num_heads, 1"
    ]
    venv_check_call(llm_venv, compile_cmd)

    # Kernel for data type=float32, BLOCK_M=64, BLOCK_DMODEL=64, BLOCK_N=64
    compile_cmd = [
        f"{llm_openai_triton_model_root}/python/triton/tools/compile.py",
        f"{openai_triton_example_root}/fmha_triton.py",
        "-n", "fused_attention_kernel",
        "-o", f"{aot_fp32_path}/fmha_kernel_d64_fp32",
        "--out-name", "fmha_d64_fp32", "-w", "4", "-ns", num_stages,
        "-s", "*fp32:16, *fp32:16, *fp32:16, *fp32:16, *fp32:16, *fp32:16, fp32, i32, i32, i32, 64, 64, 64",
        "-g", "(seq_len + 63) / 64, batch_size * num_heads, 1"
    ]
    venv_check_call(llm_venv, compile_cmd)

    # Link generated headers and create dispatchers.
    check_call(
        f"python3 {llm_openai_triton_model_root}/python/triton/tools/link.py "
        f"{aot_fp16_path}/*.h -o {aot_path}/fmha_kernel_fp16",
        shell=True)
    check_call(
        f"python3 {llm_openai_triton_model_root}/python/triton/tools/link.py "
        f"{aot_fp32_path}/*.h -o {aot_path}/fmha_kernel_fp32",
        shell=True)

    build_path = os.path.join(openai_triton_example_root, "build")
    # yapf: enable

    # make files
    make_clean_dirs(build_path)
    cmake_args = []
    try:
        import trt_test  # noqa
    except ImportError:
        pass
    else:
        trt_include_dir, trt_lib_dir = find_tensorrt(
            trt_config["new_ld_library_path"])

    if trt_include_dir:
        cmake_args.append(f"-DTRT_INCLUDE_DIR={trt_include_dir}")

    if trt_lib_dir:
        cmake_args.append(f"-DTRT_LIB_DIR={trt_lib_dir}")

    if is_trt_environment:
        cmake_args.append(f"-DCMAKE_C_FLAGS='-D_GLIBCXX_USE_CXX11_ABI=0'")

    cmake_args = " ".join(cmake_args)
    check_call(f"cd {build_path} && cmake {cmake_args} .. && make", shell=True)

    # build engine
    build_cmd = [
        f"{openai_triton_example_root}/build.py", "--num_heads=32",
        "--head_size=64", "--max_batch_size=8", "--max_seq_len=512",
        "--dtype=float16", f"--output={engine_dir}"
    ]
    venv_check_call(llm_venv, build_cmd)

    # run inference
    run_cmd = [
        f"{openai_triton_example_root}/run.py", "--num_heads=32",
        "--head_size=64", "--batch_size=8", "--seq_len=512",
        "--log_level=verbose", "--benchmark", f"--engine_dir={engine_dir}"
    ]
    venv_check_call(llm_venv, run_cmd)


# TODO[chunweiy]: Enable it later
def test_llm_openai_triton_plugingen_1gpu(openai_triton_plugingen_example_root,
                                          openai_triton_example_root,
                                          llm_openai_triton_model_root,
                                          plugin_gen_path, llm_venv,
                                          trt_config):
    # copy the triton kernel definition
    subprocess.run(
        f"cp {openai_triton_example_root}/fmha_triton.py {openai_triton_plugingen_example_root}/fmha_triton.py"
        .split(),
        check=True)

    # generate plugin
    cmd = [
        plugin_gen_path,
        "--workspace",
        "./tmp",
        "--kernel_config",
        os.path.join(openai_triton_plugingen_example_root, "kernel_config.py"),
    ]
    try:
        import trt_test  # noqa
    except ImportError:
        pass
    else:
        trt_include_dir, trt_lib_dir = find_tensorrt(
            trt_config["new_ld_library_path"])
        if trt_lib_dir is not None:
            cmd.append(f'--trt_lib_dir={trt_lib_dir}')
        if trt_include_dir is not None:
            cmd.append(f'--trt_include_dir={trt_include_dir}')

    venv_check_call(llm_venv, cmd)

    # build engine
    cmd = [
        os.path.join(openai_triton_plugingen_example_root, "build_engine.py"),
    ]
    venv_check_call(llm_venv, cmd)

    # run engine
    cmd = [
        os.path.join(openai_triton_plugingen_example_root, "run_engine.py"),
    ]
    venv_check_call(llm_venv, cmd)