Skip to content

Llama3.1-8B with Flash Attention

Flash Attention enables efficient handling of long contexts in models like Llama3.1-8B by reducing memory usage and improving throughput. This tutorial shows how to serve Llama3.1-8B with Flash Attention using RBLN vLLM and NVIDIA Triton Inference Server.

You can check out the actual commands required to compile the model and initialize triton vllm_backend on the model zoo.

Note

This tutorial assumes that you are familiar with compiling and running inference using the RBLN SDK. If you are not familiar with RBLN SDK, refer to RBLN Optimum tutorials and the API Documentation.

Prerequisites

Compile Llama3.1-8B

You need to compile the Llama3.1-8B model using optimum-rbln.

get_model.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
import os

from optimum.rbln import RBLNLlamaForCausalLM

model_id = "meta-llama/Llama-3.1-8B-Instruct"

# Compile and export
model = RBLNLlamaForCausalLM.from_pretrained(
    model_id=model_id,
    export=True,  # Export a PyTorch model to RBLN model with Optimum
    rbln_batch_size=1,  # Batch size
    rbln_max_seq_len=131_072,  # Maximum sequence length
    rbln_num_devices=1,  # Number of NPUs
    rbln_kvcache_partition_len=8_192,  # Length of KV cache partitions for flash attention
)

# Save compiled results to disk
model.save_pretrained(os.path.basename(model_id))

Note

Choose an appropriate batch size for your serving needs. Here, it is set to 1.

get_model.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
import os

from optimum.rbln import RBLNLlamaForCausalLM

model_id = "meta-llama/Llama-3.1-8B-Instruct"

# Compile and export
model = RBLNLlamaForCausalLM.from_pretrained(
    model_id=model_id,
    export=True,  # Export a PyTorch model to RBLN model with Optimum
    rbln_batch_size=1,  # Batch size
    rbln_max_seq_len=131_072,  # Maximum sequence length
    rbln_num_devices=8,  # Number of NPUs
    rbln_kvcache_partition_len=16_384,  # Length of KV cache partitions for flash attention
)

# Save compiled results to disk
model.save_pretrained(os.path.basename(model_id))

Note

Choose an appropriate batch size for your serving needs. Here, it is set to 1. The Llama3.1-8B model requires 8 RBLN NPUs, and you can find the recommended number of NPUs for each model in RBLN Optimum Multi-NPUs Supported Models.

Triton Inference Server with vLLM enabled

NVIDIA Triton Inference Server provides vLLM backend.

Step 1. Setting Up the Backend.AI Environment

Info

This step applies only when you use Backend.AI on kt cloud. If you are not using it, skip to Step 2.

  1. Start a session via Backend.AI.
  2. Select Triton Server (ngc-triton) as your environment. You can see the version of 26.07 / vllm / x86_64 / python-py3.

Step 2. Prepare NVIDIA Triton vllm_backend and Modify Model Configurations for Llama3.1-8B

A. Clone the NVIDIA Triton Inference Server vllm_backend repository: 

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

B. Place the precompiled Llama-3.1-8B-Instruct directory into the cloned vllm_backend/samples/model_repository/vllm_model/1 directory:

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

Your directory should look like the following at this point:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
+-- vllm_backend/                               # Main directory for Triton vLLM backend
|   +-- samples/                                # Application example directory
|   |   +-- model_repository/                   # Triton model repository (model_repositories)
|   |   |   +-- vllm_model/                     # Individual model for Triton serving
|   |   |   |   +-- config.pbtxt                # Triton model configuration file
|   |   |   |   +-- 1/                          # Version directory
|   |   |   |   |   +-- model.json              # Model configuration for vLLM serving
|   |   |   |   |   +-- Llama-3.1-8B-Instruct/  # rbln compiled model files
|   |   |   |   |   |   +-- decoder.rbln
|   |   |   |   |   |   +-- prefill.rbln
|   |   |   |   |   |   +-- config.json
|   |   |   |   |   |   +-- (other model files)
|   |   +-- (other example files)
|   +-- (other backend files)

Note

  1. The vLLM backend for NVIDIA Triton Server doesn't need a model.py file unlike other vision model backends. All model processing logic is pre-included in the Docker container at backends/vllm/model.py, so you only need model.json for configuration.
  2. You can either use the default config.pbtxt from the repository or create a new one using the template below. Note that input and output formats must match exactly as shown, since they're required by the vLLM backend (see Step 4: gRPC Client Inference Request).
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
    name: "vllm_model"
    backend: "vllm"

    input [
        {
            name: "text_input"
            data_type: TYPE_STRING
            dims: [ 1 ]
        },
        {
            name: "stream"
            data_type: TYPE_BOOL
            dims: [ 1 ]
        }
    ]

    output [
        {
            name: "text_output"
            data_type: TYPE_STRING
            dims: [ 1 ]
        }
    ]

    instance_group [
        {
            count: 1
            kind: KIND_MODEL
        }
    ]

C. Modify model.json

Modify vllm_backend/samples/model_repository/vllm_model/1/model.json.

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

Step 3. Run the Inference Server

Follow the section for your environment and start the inference server.

Backend.AI

Info

This section applies only when you use Backend.AI on kt cloud. If you are not using it, skip to On-premises.

Before proceeding, install the required dependencies:

1
2
3
4
5
6
7
8
$ pip install \
  --extra-index-url https://pypi.rbln.ai/simple \
  rebel-compiler==0.12.0
$ pip install \
  --extra-index-url https://wheels.vllm.ai/0.30.0/cpu \
  --extra-index-url https://pypi.rbln.ai/simple \
  --extra-index-url https://download.pytorch.org/whl/cpu \
  "vllm-rbln[kv-connectors]==0.12.0"

Start the Triton Server:

1
$ tritonserver --model-repository /PATH/TO/YOUR/vllm_backend/samples/model_repository

The following messages indicate that the server started successfully:

1
2
3
Started GRPCInferenceService at 0.0.0.0:8001
Started HTTPService at 0.0.0.0:8000
Started Metrics Service at 0.0.0.0:8002

On-premises

To access the RBLN NPU devices, RBLN Container Toolkit must be installed and the cloned vllm_backend repository should also be mounted.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
$ docker run \
    --device rebellions.ai/npu=all \
    --privileged \
    --shm-size=1g \
    --ulimit memlock=-1 \
    --ulimit stack=67108864 \
    -v /PATH/TO/YOUR/vllm_backend:/opt/tritonserver/vllm_backend \
    -p 8000:8000 \
    -p 8001:8001 \
    -p 8002:8002 \
    -ti \
    nvcr.io/nvidia/tritonserver:26.07-vllm-python-py3

Install the required dependencies in the container:

1
2
3
4
5
6
7
8
$ pip install \
  --extra-index-url https://pypi.rbln.ai/simple \
  rebel-compiler==0.12.0
$ pip install \
  --extra-index-url https://wheels.vllm.ai/0.30.0/cpu \
  --extra-index-url https://pypi.rbln.ai/simple \
  --extra-index-url https://download.pytorch.org/whl/cpu \
  "vllm-rbln[kv-connectors]==0.12.0"

Start the Triton Server in the container: 

1
$ tritonserver --model-repository /opt/tritonserver/vllm_backend/samples/model_repository

The following messages indicate that the server started successfully:

1
2
3
Started GRPCInferenceService at 0.0.0.0:8001
Started HTTPService at 0.0.0.0:8000
Started Metrics Service at 0.0.0.0:8002

Step 4. Send a Chat Completion Request

The vLLM backend of the Triton inference server defines its own model.py, which is different from the model.py defined in the Resnet50 tutorial in terms of input/output signature. The input parameter was called INPUT__0 and the output was called OUTPUT__0. But the input parameter of vLLM has the name text_input and the output is called text_output. The client should be modified accordingly. Refer to vLLM model.py for more details.

The following shows the client code for vLLM backend. This client also requires tritonclient and grpcio packages.

Note

You need to appropriately apply a chat template for proper functioning. The chat template must be formatted with system, user, and assistant roles. Additionally, you must configure the sampling parameters like this:

1
2
3
4
sampling_params = {
    "temperature": 0.0,
    "stop": ["[User]", "[System]", "[Assistant]"],  # stop tokens
}
Refer to the sample code.

simple_vllm_client.py
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
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())

If the request works properly, the output is similar to the following.

1
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!

If you need to change other sampling parameters (such as temperature, top_p, top_k, max_tokens, early_stopping), refer to VLLM's python client.

References