Transformers documentation

연속 배치

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.22.1).
Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

연속 배치

Continuous batching은 각 생성 단계마다 배치 동적 재스케줄링을 통해 GPU 활용률을 극대화합니다. 요청이 완료되면 나머지 배치 전체가 완료될 때까지 기다리지 않고 즉시 새로운 요청이 합류합니다. GPU는 항상 포화 상태를 유지하고 처리량이 높게 유지됩니다.

프로덕션 배포에는 transformers serve를 사용하세요. ContinuousBatchingManager를 기반으로 하며 OpenAI 호환 HTTP 엔드포인트를 제공합니다.

generate_batch

Continuous batching은 generate_batch()를 통해 지원됩니다. 토큰화된 프롬프트 목록을 전달하면 모든 프롬프트의 처리가 완료된 후 결과를 반환받습니다. generate_batch는 내부적으로 스케줄링을 처리하며 모든 요청이 완료될 때까지 블로킹됩니다.

서빙 및 스트리밍 사용 사례의 경우, ContinuousBatchingManager를 직접 사용하여 요청을 관리합니다.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-4B")

prompts = [
    "What's up?",
    "Name a cat breed.",
    "Write a detailed history of quantum mechanics.",
]
inputs = [tokenizer.encode(p) for p in prompts]

generation_config = GenerationConfig(
    max_new_tokens=64,
    eos_token_id=tokenizer.eos_token_id,
)

outputs = model.generate_batch(inputs=inputs, generation_config=generation_config)

for request_id, output in outputs.items():
    text = tokenizer.decode(output.generated_tokens, skip_special_tokens=True)
    print(f"[{request_id}] {text}")

ContinuousBatchingManager

ContinuousBatchingManager는 백그라운드 스레드를 실행하여 요청을 제출하고 결과를 독립적으로 가져오도록 합니다. 각 생성 단계마다 완료된 요청을 확인하고 배치에 합류할 새 요청을 스케줄링합니다. 이는 스트리밍, 실시간 서빙, 또는 요청이 도착하는 대로 제출하는 경우에 유용합니다.

continuous_batching_context_manager()를 사용하여 매니저를 안전하게 시작하고 중지합니다. 아래 예시에는 가변 길이 입력이 포함되어 있습니다. 가장 짧은 프롬프트가 완료되는 즉시 해당 프롬프트는 배치에서 빠져나가고, 더 긴 프롬프트는 생성을 계속합니다. 정적 배치(static batching)를 사용했다면 모든 프롬프트를 동일한 길이로 패딩해야 합니다. 연속 배치(continuous batching)는 완료된 프롬프트의 자원을 즉시 해제하여 다음 프롬프트의 처리를 바로 시작할 수 있게 합니다.

with model.continuous_batching_context_manager(generation_config=generation_config) as manager:
    manager.add_request(
        input_ids=tokenizer.encode("Write a detailed history of quantum mechanics."),
        request_id="long",
        max_new_tokens=512,
    )
    manager.add_request(
        input_ids=tokenizer.encode("What's up?"),
        request_id="short_0",
        max_new_tokens=32,
    )
    manager.add_request(
        input_ids=tokenizer.encode("Name a cat breed."),
        request_id="short_1",
        max_new_tokens=32,
    )

    for result in manager:
        text = tokenizer.decode(result.generated_tokens, skip_special_tokens=True)
        print(f"[{result.request_id}] {text}")

init_continuous_batching()를 호출하여 라이프사이클을 직접 관리할 수도 있습니다.

manager = model.init_continuous_batching(generation_config=generation_config)
manager.start()

# submit and retrieve requests...

매니저 종료

Manager는 백그라운드 스레드를 실행하고 분산 리소스를 보유합니다. 종료는 두 단계로 이루어지며, 진행 중인 작업에 대해 어떻게 처리할지 선택할 수 있습니다.

stop()을 호출하여 백그라운드 스레드를 중단합니다. 기본적으로 매니저는 새로운 제출을 더 이상 받지 않고, 큐에 대기 중인 요청과 활성 요청이 모두 완료된 후에 스레드가 종료됩니다.

manager.stop()

hard_stop=True를 전달하여 대기 중인 작업을 즉시 중단합니다. 큐에 대기 중이거나 활성 상태인 요청은 완료되지 않고 RuntimeError로 실패합니다.

manager.stop(hard_stop=True)

stop이 호출되면 add_request()와 add_requests()는 새로운 제출을 거부하고 경고를 기록합니다. 동일한 매니저로 또 다른 생성 세션을 실행하려면 start를 다시 호출할 수 있습니다.

destroy()를 호출하여 분산 리소스를 해제합니다. destroy는 매니저가 아직 실행 중이면 먼저 매니저를 중지하며, 이후에는 매니저를 다시 시작할 수 없습니다. 프로세스의 수명 동안 연속 배치 처리를 완료한 시점에 사용하세요.

manager.destroy()

continuous_batching_context_manager()가 이 프로세스를 처리합니다. persistent_manager=True를 전달하여 다음 세션을 위해 모델을 캐시하지 않는 한, 종료 시 stop을 호출하고 destroy를 실행합니다.

요청 추가

add_request()는 단일 요청을 제출합니다. request_id를 제공하거나 매니저가 자동으로 생성하도록 할 수 있습니다.

manager.add_request(input_ids=input_ids, request_id="my_request")

add_requests()는 배치 전체를 한 번에 제출합니다. 블록 공유가 활성화된 경우, 입력을 자동으로 정렬하여 프리픽스 캐시 히트율을 최대화합니다.

manager.add_requests(inputs=inputs)

cancel_request()로 요청을 취소합니다.

manager.cancel_request(request_id="my_request")

요청별 샘플링 매개변수

per_request_processors를 활성화하면 동일한 forward pass 내에서 각 요청별로 temperature, top_k, top_p를 독립적으로 적용할 수 있어, 요청마다 다른 샘플링 매개변수를 사용할 수 있습니다(예: 창의적이고 높은 temperature의 출력과 정확하고 낮은 temperature의 출력).

cb_config = ContinuousBatchingConfig(per_request_processors=True)

# each request gets its own sampling parameters
manager.add_request(input_ids=inputs_a, temperature=0.9, top_p=0.95)
manager.add_request(input_ids=inputs_b, temperature=0.1, top_k=10)

GenerationConfig의 각 매개변수는 런타임에 관련 logits 프로세서가 생성되도록 기본값이 아닌 값이어야 합니다. 예를 들어, 요청별 temperature 제어를 지원하려면 temperature를 None 또는 1이 아닌 값으로 설정해야 합니다. temperature가 1인 요청은 이후에도 계속 생성할 수 있습니다.

결과 가져오기

manager를 반복하여 결과가 도착하는 대로 수신합니다.

for result in manager:
    print(tokenizer.decode(result.generated_tokens, skip_special_tokens=True))

get_result()은 출력 큐에서 다음 결과를 가져옵니다. 특정 요청을 필터링하려면 request_id를 전달합니다. 큐의 다음 결과가 일치하지 않으면 다시 큐에 추가되고 메서드는 None을 반환합니다.

# next available result
result = manager.get_result()

# filter for a specific request
result = manager.get_result(request_id="my_request")

스트리밍

요청에 streaming=True를 설정한 후, request_id_iter()를 사용하여 토큰이 생성되는 동안 부분 출력을 반복합니다.

from transformers.generation.continuous_batching import RequestStatus

manager.add_request(input_ids=input_ids, request_id="streamed", streaming=True)

for chunk in manager.request_id_iter(request_id="streamed"):
    token = tokenizer.decode(chunk.generated_tokens[-1:], skip_special_tokens=True)
    print(token, end="", flush=True)
    if chunk.status == RequestStatus.FINISHED:
        break

ContinuousBatchingConfig

ContinuousBatchingConfig은 KV 캐시, 스케줄링, CUDA 그래프, 메모리 사용량 등을 제어합니다. generate_batch 또는 init_continuous_batching()에 continuous_batching_config 인수로 전달하며, GenerationConfig와 별도로 전달합니다. 두 설정은 서로 다른 것을 설명합니다. GenerationConfig는 모델 중심이며 샘플링 및 중단 매개변수를 포함하고, ContinuousBatchingConfig는 하드웨어 중심이며 메모리 및 스케줄링 매개변수를 포함합니다.

GenerationConfig에 continuous_batching_config를 설정하는 것은 비추천되며, FutureWarning을 발생시키고 v5.19에서 제거됩니다.

기본적으로 max_batch_tokens는 8192이며, 사용 가능한 GPU 메모리에 의해 제한되고 256 미만으로 설정되지 않습니다. num_blocks는 나머지 메모리를 채웁니다. 아래 표를 참고하여 적절한 기능을 선택하세요.

기능메모리처리량지연 시간
max_memory_percent / page_size✓ KV 예산을 제어합니다
max_batch_tokens↑ 더 큰 입력 버퍼✓ 더 큰 prefill 배치✓ prefill-bound일 때 TTFT
scheduler✓ 스케줄링 정책✓ TTFT
CUDA 그래프↑ 그래프 저장✓ 디스패치 오버헤드 감소✓
Async batching↑ ~2× I/O 버퍼✓ CPU/GPU 중복
컴파일↑ warmup-time만✓ 더 빠른 forward pass✓
빠른 디코딩 경로↑ 요청별 블록 테이블✓ 더 빠른 디코드 전용 단계✓
CPU 오프로딩↑ 고정된 CPU 메모리✓ 일부 재프리필을 건너뜁니다
Prefix 캐싱↓ 공유 KV 블록✓ 중복 prefill 건너뜀✓ TTFT
Paged attention↓ 조각 나지 않음✓ 동적 배치 멤버십
슬라이딩 윈도우↓ 레이어당 제한된 KV
요청별 프로세서✓ 배치별 혼합 샘플링 매개변수
max_requests_per_batch↓ caps logits 버퍼✓ 배치 크기 제한
safety_margin✓ 디코딩 지연 시간 보호
from transformers.generation import ContinuousBatchingConfig

cb_config = ContinuousBatchingConfig(
    max_memory_percent=0.8,
    page_size=256,           # KV cache page size in tokens
    scheduler_type="fifo",        # "fifo" or "prefill_first"
)

outputs = model.generate_batch(
    inputs=inputs,
    generation_config=generation_config,
    continuous_batching_config=cb_config,
)

KV 캐시 페이지 크기

page_size는 한 레이어의 풀 어텐션 페이지가 KV 토큰을 얼마나 담을지 설정합니다. 최소 4 이상이어야 하며, 그렇지 않으면 초기화 시 ValueError가 발생합니다. 기본값 256은 대부분의 워크로드에 적합합니다.

block_size 매개변수는 더 이상 사용되지 않으며 page_size로 대체되었습니다.

페이지 크기가 크면 관리 작업은 줄지만, 시퀀스가 페이지 중간에서 끝나면 공간이 낭비됩니다. 페이지 크기가 작으면 공간 낭비는 적지만 관리 작업이 더 많이 필요합니다. 일반적으로 page_size를 조정할 필요가 없습니다. 짧은 요청이 많은 워크로드라면 더 작은 값을, 매우 긴 요청이 몇 개 있는 워크로드라면 더 큰 값을 사용해 보세요. 페이지, 블록, 청크가 어떻게 결합되는지 확인하려면 메모리 관리를 참조하세요.

cb_config = ContinuousBatchingConfig(page_size=128)

Prefill 배치 크기

max_batch_tokens는 단일 forward pass의 토큰 예산, 즉 모델이 한 번에 처리할 수 있는 최대 쿼리 토큰 수를 설정합니다. 예산이 클수록 각 prefill에 더 많은 프롬프트 토큰이 포함되므로, 프롬프트 중심 워크로드에서 prefill 처리량이 높아지고 첫 토큰까지의 시간이 줄어듭니다. 스케줄러는 예산을 초과하는 프롬프트를 여러 스텝에 걸쳐 분할합니다. 크기가 초과된 프롬프트가 처리되는 방식은 chunked prefill을 참고하세요.

기본적으로 max_batch_tokens의 값은 8192입니다. 이 값은 사용 가능한 GPU 메모리에 의해 제한되며, 배치별 입력 텐서가 KV 캐시를 밀어내지 않도록 합니다. 또한 이 값은 256 미만으로 내려가지 않습니다. 여유 메모리가 적은 GPU에서는 이 제한으로 인해 값이 8192보다 낮아질 수 있습니다.

max_batch_tokens와 num_blocks는 동일한 메모리 예산에서 할당되므로 서로 트레이드오프 관계에 있습니다. 토큰 예산을 크게 설정하면 더 큰 입력 버퍼가 할당되고, 동시 또는 긴 요청을 위한 KV 캐시 블록은 줄어듭니다. 사용 가능한 메모리가 충분할 때 prefill 처리량을 높이기 위해 max_batch_tokens를 높이고, 더 많은 KV 블록에 메모리를 확보하거나 OOM 오류를 방지하기 위해 낮추세요.

cb_config = ContinuousBatchingConfig(max_batch_tokens=16384)

배치 및 스케줄링 한도

max_requests_per_batch는 단일 forward pass에서 실행되는 요청 수의 상한을 설정합니다. 모델은 배치 토큰 수와 어휘 크기로 크기가 결정되는 logits 텐서를 생성하며, 샘플링을 위해 fp32로 캐스팅되어 메모리 사용량이 두 배가 됩니다. 각 요청은 다음 토큰 예측 하나만 필요하므로, 상한을 낮추면 이 텐서의 크기를 줄이고 대규모 어휘를 가진 prefill 중심 배치에서 out-of-memory 오류를 방지할 수 있습니다. 기본값은 제출된 프롬프트 수이며, 1024로 폴백된 후 max_batch_tokens와 num_blocks에 맞춰 상한이 적용됩니다.

cb_config = ContinuousBatchingConfig(max_requests_per_batch=256)

safety_margin은 활성 요청을 위해 사용 가능한 캐시 용량의 일부를 예약합니다. 사용 가능한 용량이 safety_margin(캐시 전체 크기의 비율) 미만으로 떨어지면, 스케줄러는 새로운 prefill의 수락을 중단하고 이미 진행 중인 디코딩 요청만 계속 처리합니다. 이는 새로운 요청을 시작하는 것보다 활성 작업을 완료하는 것을 우선시하여 디코딩 지연 시간을 보호하고 offloading을 지연시킵니다. 값은 0과 1 사이여야 하며, 0은 마진을 비활성화합니다.

cb_config = ContinuousBatchingConfig(safety_margin=0.15)

기본값은 스케줄러에 따라 다릅니다. FIFO 스케줄러의 경우 0.15, prefill_first의 경우 0.0입니다. 마진이 스케줄링과 어떻게 상호작용하는지는 Admission을 참고하세요.

로그 확률

return_logprobs=True가 설정된 경우 ContinuousBatchingConfig는 각 생성된 토큰의 로그 확률을 반환합니다. 이는 RL에서 로그 확률이 일부 학습 루프의 입력으로 사용되는 경우에 유용합니다.

cb_config = ContinuousBatchingConfig(return_logprobs=True)

# generate_batch()

for request_id, output in outputs.items():
    for token_id, log_prob in zip(output.generated_tokens, output.logprobs):
        token = tokenizer.decode([token_id])
        print(f"{token} | logprob: {log_prob}")

CUDA 그래프

CUDA graphs는 GPU 실행 그래프를 한 번 기록한 후 동일한 shape를 가진 배치에 대해 이를 재실행함으로써 CPU 디스패치 오버헤드를 제거합니다. use_cuda_graph=True를 명시적으로 설정하여 활성화합니다.

cb_config = ContinuousBatchingConfig(use_cuda_graph=True)

활성화되면, 매니저는 쿼리 및 KV 길이를 고정된 간격으로 패딩하여 모양이 반복되고 그래프가 재사용되도록 합니다. q_padding_interval_size와 kv_padding_interval_size의 값이 작을수록 패딩으로 인한 계산 낭비가 줄어들지만, 그래프가 기록하고 저장해야 하는 고유한 모양의 수가 증가하여 메모리 비용이 더 많이 듭니다.

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    q_padding_interval_size=64,
    kv_padding_interval_size=16384,
)

Async batching

Async batching은 다음 배치의 CPU 스케줄링과 현재 배치의 GPU 연산을 겹치게 합니다. CUDA graphs가 필요하며, 입력 텐서에 사용되는 VRAM이 대략 두 배로 증가합니다.

cb_config = ContinuousBatchingConfig(
    use_cuda_graph=True,
    use_async_batching=True,
)

컴파일

default_compile_level은 모델의 forward pass에 torch.compile을 적용합니다. 컴파일은 일회성 warmup 비용을 지불하는 대신 생성 과정에서의 forward pass 속도를 높입니다. 레벨이 높을수록 더 공격적인 최적화가 수행되어 처리량이 개선되지만 warmup 시간이 길어집니다. 테스트 및 벤치마크 반복 작업에는 레벨을 낮게 유지하고, warmup 비용이 수많은 요청에 걸쳐 상쇄되는 장기 실행 서빙 워크로드에서는 레벨을 높이는 것이 좋습니다.

레벨은 0에서 3까지의 범위를 가집니다. 레벨 0이 기본값이며, 컴파일을 완전히 건너뜁니다.

레벨modedynamic트레이드오프
0——컴파일 없음(기본값), 가장 빠른 시작
1defaultTrue적당한 속도 향상, 짧은 워밍업
2max-autotune-no-cudagraphsTrue더 큰 속도 향상, 더 긴 워밍업
3max-autotune-no-cudagraphsFalse최고의 처리량, 가장 긴 워밍업
cb_config = ContinuousBatchingConfig(default_compile_level=1)

이 레벨은 varlen 및 decode 실행 경로에 기본 CompileConfig를 제공합니다. 명시적인 설정이 없는 경로에만 적용되므로, varlen_compile_config와 decode_compile_config가 설정된 경우 우선적으로 적용됩니다. FlashAttention 환경에서는 max_length_k가 빈번한 재컴파일을 유발하여 varlen 경로가 컴파일을 건너뛰므로, 이 경우 해당 레벨은 decode 경로에만 영향을 미칩니다.

빠른 디코딩 경로

배치에 디코딩 요청만 포함되는 경우(시퀀스당 쿼리 토큰 1개)에는 가변 길이 커널 대신 flash_attn_with_kvcache 커널로 디스패치할 수 있습니다. 이 방식은 수동 업데이트를 거치지 않고 블록 테이블을 통해 페이지드 KV 캐시를 인플레이스로 읽기 및 쓰기 때문에 varlen 경로보다 빠릅니다. 커널 수준의 세부 사항은 Paged attention을 참고하세요.

Fast path는 max_blocks_per_request에 의해 크기가 결정되며, 이는 요청별 블록 테이블의 차원을 설정합니다. 기본적으로 이 값은 자동으로 추론됩니다. 매니저에 max_prompt_length와 max_generated_length가 설정되어 있으면, 블록 테이블은 최대 시퀀스 길도에 맞게 크기가 조정됩니다. 그렇지 않으면 폴백 기본값(요청당 32개 블록)이 사용됩니다.

max_blocks_per_request을 특정 값으로 설정하여 블록 테이블 크기를 명시적으로 지정할 수 있습니다. 이는 요청당 최대 시퀀스 길이를 알고 있고 블록 테이블의 메모리 비용을 제한하려는 경우에 유용합니다.

cb_config = ContinuousBatchingConfig(max_blocks_per_request=64)

max_blocks_per_request=0을 설정하면 fast path를 비활성화하고 모든 batch를 varlen kernel을 통해 처리하도록 강제합니다. 이는 기본값 이전의 동작을 복원하며, attention 구현에서 fast path를 사용할 수 없는 경우에 유용합니다(기본 kernel을 사용할 수 없는 경우 매니저가 자동으로 비활성화하기도 합니다).

cb_config = ContinuousBatchingConfig(max_blocks_per_request=0)

fast path는 두 가지의 디바이스와 어텐션 구현 조합에서 사용할 수 있는 flash_attn_with_kvcache 커널에 의존합니다.

장치attn_implementation
CUDAflash_attention_2 또는 flash_attention_3
XPUflash_attention_2

다른 조합이거나 커널을 가져올 수 없는 경우, 매니저는 varlen 경로로 폴백합니다. max_blocks_per_request를 명시적으로 설정한 경우에만 경고가 기록됩니다.

Full attention은 kernel block tables를 지원하지만 sliding attention은 지원하지 않습니다. 모델에 sliding attention 레이어가 하나라도 있는 경우, full attention을 함께 사용하는 하이브리드 모델도 포함하여, 초기화 시 max_blocks_per_request가 0으로 강제 설정되며, 사용자가 설정한 0이 아닌 값은 무시됩니다. 이후 매니저는 모델의 일부 attention 유형이 kernel block tables를 지원하지 않아 decode fast path를 사용할 수 없다는 내용을 로그로 기록합니다.

CPU 오프로딩

CPU offloading은 GPU KV 캐시가 가득 찼을 때, evicted된 KV 캐시 블록을 사전에 할당된 pinned CPU 버퍼로 복사합니다. 캐시 공간이 확보되면, 매니저는 해당 블록을 다시 GPU로 복사하고 프롬프트와 생성된 토큰을 다시 계산하지 않고 요청을 재개합니다.

cpu_offload_space를 GiB 단위의 CPU 스왑 공간으로 설정합니다. 기본값인 0.0은 CPU 오프로딩을 비활성화합니다.

cb_config = ContinuousBatchingConfig(cpu_offload_space=8.0)

기본적으로 cpu_offload_space_safety_threshold=0.8은 psutil이 설치된 경우 요청된 공간을 사용 가능한 시스템 RAM의 80%로 제한합니다. cpu_offload_space=None으로 설정하면 안전 임계값에 따라 스왑 풀 크기를 결정합니다.

Tensor parallel 타임아웃

Tensor parallelism 환경에서 매니저는 각 rank 간에 요청 제출, 취소 및 종료 작업을 조정하기 위해 CPU 통신 그룹을 생성합니다. cpu_group_timeout은 이 그룹에서 collective 연산이 프로세스가 충돌하기 전에 차단될 수 있는 최대 시간을 제한합니다. 한 rank가 멈추면, 이 타임아웃이 다른 rank들이 영원히 대기하는 것을 방지합니다.

collective를 드물게 발행하는 워크로드에 대해 더 긴 타임아웃을 설정하거나, None을 전달하여 이를 비활성화합니다.

cb_config = ContinuousBatchingConfig(cpu_group_timeout=600.0)

Prefix 캐싱

여러 요청이 시스템 프롬프트와 같은 공통 접두사를 공유하는 경우, 매니저는 이를 다시 계산하는 대신 KV 캐시 블록을 재사용합니다. allow_block_sharing의 기본값은 True입니다. 접두사 공유는 모든 어텐션 타입 할당기가 공유를 지원하는 경우에만 활성화되므로, 완전한 어텐션 전용 모델에서만 사용됩니다. 순수 슬라이딩 모델과 하이브리드 모델은 allow_block_sharing가 True인 경우에도 이를 비활성화합니다.

cb_config = ContinuousBatchingConfig(
    allow_block_sharing=True,  # default
)

어텐션 백엔드 선택

Continuous batching은 FlashAttention, SDPA, "paged|eager"를 지원합니다. 모델을 로드할 때 attn_implementation을 설정하세요.

백엔드attn_implementation요구 사항
FlashAttention 2“flash_attention_2”flash-attn 또는 kernels 패키지
FlashAttention 3“flash_attention_3”flash-attn-3 또는 kernels 패키지
SDPA (PyTorch 네이티브)“sdpa”None
Eager“paged|eager”None
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-4B",
    attn_implementation="flash_attention_2",
    device_map="auto",
    dtype=torch.bfloat16,
)

Continuous batching은 FlashAttention을 사용할 때 eager나 SDPA보다 더 빠르게 실행되며, 이는 주로 FlashAttention이 attention mask를 필요로 하지 않기 때문입니다. 모델이 FlashAttention을 지원하는 경우, "eager" 또는 "sdpa"로 로드된 모델은 FlashAttention 3으로 전환되며, FlashAttention 3이 설치되어 있지 않은 경우 FlashAttention 2로 전환되고 경고가 기록됩니다. 두 버전 모두 설치되어 있지 않으면 백엔드는 변경되지 않습니다.

테스트와 같이 eager 또는 SDPA를 유지하려면 ContinuousBatchingConfig에서 auto_switch_to_flash=False를 설정하세요. 백엔드가 "eager"로 유지되는 경우, continuous batching은 대신 "paged|eager"를 실행합니다.

Tensor 병렬성

단일 GPU에 탑재할 수 없는 모델의 경우, 텐서 병렬을 사용하여 가중치를 여러 디바이스에 샤딩합니다. DistributedConfig(tp_size=N)으로 디바이스 수를 설정합니다. 연속 배치 처리는 모델에서 텐서 병렬 크기를 읽어 샤드별 페이지드 KV 캐시 크기를 결정합니다. 지원되는 아키텍처 목록과 샤딩 작동 방식은 Tensor parallelism을 참고하세요.

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, DistributedConfig
from transformers.generation import ContinuousBatchingConfig, GenerationConfig

distributed_config = DistributedConfig(tp_size=4)
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen3-32B",
    attn_implementation="flash_attention_2",
    distributed_config=distributed_config,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-32B")

inputs = [tokenizer.encode(p) for p in ["What's up?", "Name a cat breed."]]
generation_config = GenerationConfig(max_new_tokens=64, eos_token_id=tokenizer.eos_token_id)

outputs = model.generate_batch(inputs=inputs, generation_config=generation_config)

torchrun로 스크립트를 실행하고, --nproc-per-node를 샤딩할 GPU 수로 설정합니다.

torchrun --nproc-per-node 4 cb_tp.py

Tensor parallel size은 모델의 num_key_value_heads를 나누어야 합니다(모델 설정을 확인하세요). 그렇지 않으면 paged cache가 시작 시 오류를 발생시키므로, 적절한 --nproc-per-node를 선택하세요.

device_map을 distributed_config와 함께 설정하지 마세요. device_map은 모듈 전체를 특정 GPU에 배치하는 반면, 텐서 병렬성은 동일한 매개변수를 모든 GPU에 걸쳐 분할하기 때문에 두 설정은 충돌합니다.

슬라이딩 윈도우 어텐션

슬라이딩 윈도우 어텐션(Mistral, Gemma 2)을 사용하는 모델은 연속 배치 처리와 함께 작동합니다. 슬라이딩 레이어는 SlidingAttentionCacheAllocator를 사용하며, 이는 소수의 페이지를 재사용하여 최근 토큰의 고정된 윈도우만 유지합니다. 어텐션 유형은 ContinuousBatchingConfig가 아니라 모델의 layer_types 또는 sliding_window에서 가져옵니다.

from transformers import AutoConfig, AutoModelForCausalLM

config = AutoConfig.from_pretrained("google/gemma-2-2b")
config.sliding_window = 4096

model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-2b",
    config=config,
    attn_implementation="sdpa",
    device_map="auto",
    dtype=torch.bfloat16,
)

Prefix caching와 decode fast path는 sliding attention이 존재하는 경우, full attention 레이어를 함께 가진 하이브리드 모델을 포함하여 모두 비활성화됩니다.

다음 단계

Update on GitHub