Flash Attention을 이용한 Llama3.1-8B 서빙¶

$ git clone https://github.com/triton-inference-server/vllm_backend.git -b r26.07

$ cp -R /PATH/TO/YOUR/Llama-3.1-8B-Instruct \
        /PATH/TO/YOUR/CLONED/vllm_backend/samples/model_repository/vllm_model/1/

+-- vllm_backend/                               # Triton vLLM 백엔드 메인 디렉토리
|   +-- samples/                                # 애플리케이션 예제 디렉토리
|   |   +-- model_repository/                   # 모델 저장소 (model_repositories)
|   |   |   +-- vllm_model/                     # Triton 으로 서빙할 개별 모델
|   |   |   |   +-- config.pbtxt                # Triton 모델 설정 파일
|   |   |   |   +-- 1/                          # 버전 디렉토리
|   |   |   |   |   +-- model.json              # vLLM 서빙위한 모델 설정
|   |   |   |   |   +-- Llama-3.1-8B-Instruct/  # 컴파일된 모델 파일들(rbln files)
|   |   |   |   |   |   +-- decoder.rbln
|   |   |   |   |   |   +-- prefill.rbln
|   |   |   |   |   |   +-- config.json
|   |   |   |   |   |   +-- (기타 모델 파일들)
|   |   +-- (기타 예제 파일들)
|   +-- (기타 백엔드 파일들)

{
    "model": "/ABSOLUTE/PATH/TO/Llama-3.1-8B-Instruct"
}

import asyncio
import json

import numpy as np
import tritonclient.grpc.aio as grpcclient


# Define a simple chat message class
class ChatMessage:
    def __init__(self, role, content):
        self.role = role
        self.content = content


# Apply a simple chat template to the messages
def apply_chat_template(messages):
    lines = []
    system_msg = ChatMessage(role="system", content="You are a helpful assistant.")
    for msg in [system_msg, *messages, ChatMessage(role="assistant", content="")]:
        lines.append(f"[{msg.role.capitalize()}]\n{msg.content}")
    return "\n".join(lines)


async def try_request():
    url = (
        "<host and port number of the triton inference server>"  # e.g. "localhost:8001"
    )
    client = grpcclient.InferenceServerClient(url=url, verbose=False)

    model_name = "vllm_model"

    def create_request(messages, request_id):
        prompt = apply_chat_template(messages)
        print(f"prompt:\n{prompt}\n---")  # print prompt

        input = grpcclient.InferInput("text_input", [1], "BYTES")
        prompt_data = np.array([prompt.encode("utf-8")])
        input.set_data_from_numpy(prompt_data)

        stream_setting = grpcclient.InferInput("stream", [1], "BOOL")
        stream_setting.set_data_from_numpy(np.array([True]))

        sampling_params = {
            "temperature": 0.0,
            "stop": ["[User]", "[System]", "[Assistant]"],  # add stop tokens
        }

        sampling_parameters = grpcclient.InferInput("sampling_parameters", [1], "BYTES")
        sampling_parameters.set_data_from_numpy(
            np.array([json.dumps(sampling_params).encode("utf-8")], dtype=object)
        )

        inputs = [input, stream_setting, sampling_parameters]
        output = grpcclient.InferRequestedOutput("text_output")
        outputs = [output]

        return {
            "model_name": model_name,
            "inputs": inputs,
            "outputs": outputs,
            "request_id": request_id,
        }

    messages = [
        ChatMessage(
            role="user", content="What is the first letter of English alphabets?"
        )
    ]
    request_id = "req-0"

    async def requests_gen():
        yield create_request(messages, request_id)

    response_stream = client.stream_infer(requests_gen())

    prompt = apply_chat_template(messages)
    is_first_response = True

    async for response in response_stream:
        result, error = response
        if error:
            print("Error occurred!")
        else:
            output = result.as_numpy("text_output")
            for i in output:
                decoded = i.decode("utf-8")
                if is_first_response:
                    if decoded.startswith(prompt):
                        decoded = decoded[len(prompt) :]
                    is_first_response = False
                print(decoded, end="", flush=True)
    print("\n")  # end of stream


asyncio.run(try_request())

The first letter of the English alphabet is 'A'. Would you like to know more about the English alphabet? I can help you with that!