피스와이즈 CUDA 그래프

피스와이즈 CUDA 그래프 (Piecewise CUDA Graph)

Piecewise CUDA Graph (PCG)는 모델의 계산 그래프를 "분할 지점"(예: MoE 디스패치 연산)에서 여러 조각(대략 레이어당 하나)으로 나눠, 반복마다 토큰 수가 변하는 extend/prefill에서도 커널 런치 오버헤드를 제거하면서 동적 형태(dynamic shapes)를 지원합니다. 기본적으로 활성화됩니다.

출처: 문서

본문

동기 (Motivation)

표준 CUDA 그래프는 전체 모델 forward pass를 단일 그래프로 캡처합니다. 이는 decode(고정 배치 크기)에는 잘 맞지만, 토큰 수가 반복마다 변하는 extend/prefill에는 맞지 않아요.

Piecewise CUDA Graph (PCG)는 모델의 계산 그래프를 "분할 지점"(split points, 예: MoE dispatch ops)에서 여러 조각(대략 레이어당 하나)으로 나눠 해결합니다. 각 조각은 사전 정의된 토큰 길이 집합에 대해 별도 CUDA 그래프로 캡처됩니다. 런타임에 입력은 가장 가까운 캡처 크기로 패딩되고 각 조각이 재생됩니다. 이는 동적 형태를 지원하면서 prefill/extend의 커널 런치 오버헤드를 제거합니다.

PCG는 기본적으로 활성화됩니다. 끄려면 --cuda-graph-backend-prefill=disabled를 전달하세요.

사용법 (Usage)

PCG는 지원되는 구성에서 기본 활성화됩니다. 추가 플래그가 필요 없어요:

python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct

PCG 비활성화 (Disable PCG)

python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --cuda-graph-backend-prefill=disabled

커스텀 캡처 크기 (Custom capture sizes)

python3 -m sglang.launch_server \
    --model-path meta-llama/Llama-3.1-8B-Instruct \
    --cuda-graph-max-bs-prefill 2048

서버 인자 (Server Args)

Argument	Default	Description
--cuda-graph-backend-prefill	None (auto)	prefill 단계의 백엔드. 선택지: full, breakable, tc_piecewise, disabled. disabled를 전달하면 extend/prefill에서 PCG를 끄고, tc_piecewise는 모든 자동 비활성 검사를 건너뛰고 강제로 켭니다(테스트 전용).
--cuda-graph-max-bs-prefill	None (auto)	캡처할 최대 토큰 수. 기본은 chunked_prefill_size(비-MLA) 또는 2048(MLA).
--cuda-graph-bs-prefill	None (auto)	캡처할 토큰 길이의 명시적 목록. 설정하지 않으면 자동 생성.
--cuda-graph-tc-compiler	"eager"	캡처된 서브그래프의 컴파일러 백엔드. 선택지: eager, inductor.

버그 리포트 (Bug Report)

PCG는 기본 활성화지만 여전히 실험 단계입니다. PCG는 torch.compile로 모델 forward pass를 추적하므로, 대부분의 버그는 torch compile 추적 실패(예: 추적 불가 op, 동적 제어 흐름, 그래프 중단)에서 발생합니다. PCG 관련 문제가 발생하면 시작 명령에 --cuda-graph-backend-prefill=disabled를 추가해 비활성화하고 GitHub Issues에 버그를 보고해 주세요. 이 기능을 개선하는 데 큰 도움이 됩니다.

사용자용 (For Users)

서버 시작 중 다음과 같은 오류 메시지가 보이면 PCG 버그입니다:

Piecewise CUDA Graph is enabled by default as an experimental feature.
To work around this error, add --cuda-graph-backend-prefill=disabled to your launch command.
Please report this issue at https://github.com/sgl-project/sglang/issues/new/choose

이를 우회하려면 시작 명령에 --cuda-graph-backend-prefill=disabled를 추가하세요. 버그 리포트를 작성할 때는 다음을 포함해 주세요:

1. 전체 오류 트레이스백
2. 모델명과 양자화 방법
3. 모든 인자가 있는 시작 명령
4. GPU 유형과 드라이버 버전

개발자용 (For Developers)

PCG는 torch.compile로 모델 forward pass를 추적하므로, 새로 개발된 CUDA 커널(JIT 커널과 sgl-kernels 모두)은 대개 torch.compile과 기본적으로 호환되지 않습니다. 추적은 커널 내부의 JIT 컴파일, 파일 I/O, 동적 모듈 로딩 같은 추적 불가 연산에서 실패합니다.

커널을 PCG와 호환되게 만들려면 sglang.srt.utils.custom_op의 register_custom_op로 커스텀 op로 등록해야 합니다. 이는 커널을 컴파일된 그래프에서 불투명한 노드로 감싸 torch.compile이 그 안을 추적하지 않게 합니다.

예제 사용법 (JIT 커널):

from sglang.srt.utils.custom_op import register_custom_op

# Inplace operator (no return value)
@register_custom_op(mutates_args=["output_q", "output_s"])
def per_token_group_quant_8bit(
    input: torch.Tensor,
    output_q: torch.Tensor,
    output_s: torch.Tensor,
) -> None:
    # kernel implementation ...

예제 사용법 (출력이 있는 연산자):

# out_shape indicates which argument has the same shape as the output
@register_custom_op(mutates_args=["x"], out_shape=0)
def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
    return x.add_(y)

외부 라이브러리 함수(예: FlashInfer 커널)를 감싸려면 register_custom_op_from_extern을 대신 사용하세요. 전체 API 문서는 python/sglang/srt/utils/custom_op.py를 참고하세요.

동작 방식 (How it works)

Torch compile 백엔드

PCG는 커스텀 백엔드(SGLangBackend)로 torch.compile을 사용해 모델 forward pass를 분할·컴파일합니다. 흐름은:

model.forward wrapper
→ torch.compile(..., backend=SGLangBackend)
→ FX graph
→ split_graph() at registered split ops
→ split_gm (top-level graph that chains the pieces)
→ replace capturable submodules with CUDAPiecewiseBackend
→ runtime dispatch: eager split ops + per-piece capture/replay

• 설치 (Install): install_torch_compiled()가 model.forward를 래퍼 함수로 교체합니다. is_in_piecewise_cuda_graph()가 True를 반환하면 래퍼는 컴파일된 callable로 디스패치하고, 그 외에는 원래 forward로 폴백합니다. 이 경로를 통한 첫 호출이 Dynamo 트레이싱과 그래프 컴파일을 트리거합니다. CUDA 그래프 재생은 캡처 단계가 끝난 후에만 발생합니다.

• 분할 (Split): torch.compile이 모델을 추적할 때 SGLangBackend가 FX 그래프를 받아 split_graph()를 호출합니다. CompilationConfig.split_ops에 나열된 연산은 분할 지점으로 취급되어 각 지점에서 그래프가 잘립니다. 이 split-op 서브모듈은 런타임에서 eager로 실행되도록 남겨지고, 주변 서브모듈은 CUDAPiecewiseBackend로 컴파일·감싸집니다. 결과는 캡처 가능한 서브그래프와 eager split-op 서브모듈이 교차하는 submod_0, submod_1, … 자식이 있는 최상위 "스티칭 그래프"(split_gm)입니다.

• 교체 (Replace): PiecewiseCompileInterpreter가 split_gm의 각 캡처 가능 서브모듈을 반복해 일반(동적) 형태로 컴파일하고, 그 자리를 CUDAPiecewiseBackend 인스턴스로 교체합니다. Split-op 서브모듈(예: attention, all-reduce)은 그대로 두고 런타임에서 eager로 실행합니다.

• 디스패치 (Dispatch): 런타임에서 split_gm을 호출하면 스티칭 그래프가 실행되어 각 서브모듈을 순서대로 호출합니다. Split-op 서브모듈은 eager로 실행됩니다. 각 CUDAPiecewiseBackend 서브모듈은 세 단계를 거칩니다:

  • Compile warmup — 일반 형태 컴파일 경로를 실행.
  • Capture — 각 캡처 크기에 대해 warmup 통과 후 CUDA 그래프를 기록.
  • Steady-state replay — 각 forward pass에 대해 캡처된 CUDA 그래프를 재생.

피스와이즈 cuda graph 러너

PiecewiseCudaGraphRunner는 세 단계를 통해 전체 수명주기를 조정합니다:

• Compile — dummy forward pass로 JIT 커널을 워밍업한 뒤 모델을 torch.compile로 감싸 Dynamo 트레이싱을 트리거해 FX 그래프를 분할하고 각 서브그래프 조각에 CUDAPiecewiseBackend 인스턴스를 생성.
• Capture — 캡처 크기를 역순(큰 것부터)으로 반복. 각 크기에 대해 forward pass를 두 번 실행(워밍업 1회, CUDA 그래프 캡처 1회).
• Replay — 런타임에서 이진 검색으로 실제 토큰 수 >= 캡처 크기 중 가장 작은 것을 찾아 입력을 정적 버퍼에 zero-padding으로 복사하고 캡처된 CUDA 그래프를 재생하며, 출력을 실제 토큰 수로 슬라이스.

메모리 최적화 (Memory optimization)

PCG의 메모리 비용은 두 부분에서 옵니다: torch memory allocator와 non-torch memory.

torch 메모리 할당자 오버헤드는 몇 가지 최적화 덕분에 사소합니다. 전역 공유 메모리 풀이 모든 CUDA 그래프 러너와 캡처 크기에 걸쳐 재사용되고, 캡처가 역순(큰 것부터 작은 것)으로 수행되어 작은 그래프가 큰 그래프가 할당한 메모리를 재사용하며, 마지막 서브그래프의 출력 텐서는 최대 메모리 재사용을 위해 약한 참조(weak references)로 저장됩니다.

주요 메모리 오버헤드는 non-torch 메모리에서 옵니다. CUDA 그래프 객체 자체가 기록된 커널 런치 파라미터와 내부 상태를 저장하려면 GPU 메모리가 필요합니다. 이 오버헤드는 캡처된 크기 수에 따라 확장되므로, piecewise_cuda_graph_max_tokens는 기본적으로 보수적으로 상한을 둡니다.

형태 구성 (Shape configuration)

Piecewise CUDA 그래프는 토큰 수 집합에 대한 그래프를 미리 캡처합니다. 런타임에 실제 토큰 수는 가장 가까운 캡처 크기로 올림(이진 검색)되고 해당 그래프가 재생됩니다. 토큰 수가 가장 큰 캡처 크기를 초과하면 런타임은 일반(비그래프) forward 경로로 폴백합니다.

기본 캡처 스케줄은 증가하는 세분성으로 자동 생성됩니다:

Token range	Step size
4 – 32	4
48 – 256	16
288 – 512	32
576 – 1024	64
1280 – 4096	256
4096+	512

자동 생성 스케줄의 경우 크기는 --cuda-graph-max-bs-prefill로 상한이 정해집니다. 기본 상한은 비-MLA 모델의 chunked_prefill_size, MLA 백엔드 모델의 2048입니다. --max-total-tokens가 설정되면 상한은 이를 넘지 않도록 더 제한됩니다. 또한 Llama-2 모델은 임시 해결책으로 4096 토큰에서 자동 상한이 정해집니다.

호환성 (Compatibility)

PCG는 다음 시나리오에서 자동 비활성화됩니다. 우리는 호환성 확장을 적극적으로 작업 중이며, 이들 중 많은 지원이 곧 제공될 예정입니다.

• 비활성 모델 아키텍처 (예: DeepseekV32ForCausalLM)
• 추측 디코딩 (Speculative decoding)
• DP attention
• 파이프라인 병렬화 (pp_size > 1)
• 비-CUDA 하드웨어 (AMD ROCm, Ascend NPU)
• MoE A2A 백엔드
• LoRA
• 멀티모달 / VLM 모델
• DLLM (diffusion LLM)
• 결정적 추론 (Deterministic inference)
• PD disaggregation
• Expert distribution recorder / EPLB

모든 자동 비활성 검사를 건너뛰려면 --cuda-graph-backend-prefill=tc_piecewise를 사용하세요(테스트/디버깅 전용).

코드 참조 (Code Reference)

File	Description
python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py	백엔드 구현: compile, capture, replay
python/sglang/srt/compilation/compile.py	install_torch_compiled 트램폴린
python/sglang/srt/compilation/backend.py	SGLangBackend, 그래프 분할, piecewise 컴파일
python/sglang/srt/compilation/cuda_piecewise_backend.py	서브그래프별 CUDA 그래프 캡처/재생
python/sglang/srt/compilation/piecewise_context_manager.py	전역 컨텍스트 플래그와 ForwardContext
python/sglang/srt/compilation/compilation_config.py	캡처 크기, split ops, 컴파일러 설정
python/sglang/srt/utils/custom_op.py	torch.compile 호환용 register_custom_op
python/sglang/srt/server_args.py	서버 인자와 자동 비활성 로직

더 알아보기 (Learn more)

• Breakable CUDA Graph