콘텐츠로 이동

Attention 모드

개요

vLLM RBLN은 Naive Attention과 Flash Attention 두 가지 모드를 지원합니다.

Attention 모드는 모델 컴파일 시 rbln_attn_impl 및 rbln_kvcache_partition_len 파라미터를 통해 지정할 수 있습니다. 추론 시에는 컴파일된 모델의 rbln_config.json 파일에 포함된 attn_impl 필드를 기반으로 자동 적용됩니다.

따라서 사용자는 컴파일 시 한 번만 Attention 모드를 지정하면 됩니다.

Naive Attention

Naive Attention은 별도의 최적화가 적용되지 않은 기본 Attention 방식입니다.

Naive Attention을 사용해 모델을 컴파일하고 실행하려면 rbln_attn_impl="eager"를 설정합니다.

rbln_attn_impl의 기본값은 "eager"이므로, 이 파라미터를 생략해도 자동으로 적용됩니다. 또한 Naive Attention에서는 rbln_kvcache_partition_len을 설정할 필요가 없습니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
from optimum.rbln import RBLNLlamaForCausalLM

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"

# Compile and export
model = RBLNLlamaForCausalLM.from_pretrained(
    model_id=model_id,
    export=True,
    rbln_attn_impl="eager",
    rbln_batch_size=4,
    rbln_max_seq_len=8192,
    rbln_tensor_parallel_size=4,
)

# Save compiled results to disk
model.save_pretrained("rbln-Llama-3-8B-Instruct")

자세한 예시는 Llama3 8B 튜토리얼을 참고하세요.

Flash Attention

Flash Attention은 메모리 효율성과 처리량을 개선하여 긴 컨텍스트를 다루는 모델의 성능을 향상시킵니다.

Flash Attention 모드는 컴파일 시 rbln_attn_impl="flash_attn"을 설정하고 rbln_kvcache_partition_len 파라미터를 지정하면 활성화됩니다. 또한 rbln_kvcache_partition_len이 지정되고 rbln_max_seq_len과 다른 값으로 설정된 경우, rbln_attn_impl="flash_attn"을 명시적으로 지정하지 않아도 flash_attn 모드가 자동으로 적용됩니다.

Flash Attention에서 최상의 성능을 얻으려면 RBLN Model Zoo에 설정된 매개변수들을 따르는 것을 추천합니다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
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,
    rbln_attn_impl="flash_attn",
    rbln_batch_size=1,
    rbln_max_seq_len=131_072,
    rbln_tensor_parallel_size=8,
    rbln_kvcache_partition_len=16_384,
)


# Save compiled results to disk
model.save_pretrained("rbln-Llama-3-1-8B-Instruct")

자세한 예시는 Llama3.1 8B 튜토리얼을 참고하세요.