본문으로 건너뛰기

[SGLang] Piecewise CUDA Graph: 분할 그래프 컴파일 전략

들어가며

일반 CUDA Graph는 Decode 단계(고정 형상)에 적합하지만, Prefill 단계는 입력 길이가 가변적이라 적용이 어렵다. SGLang은 이 문제를 "Piecewise CUDA Graph"로 해결한다. 모델을 여러 조각(piece)으로 분할하고, 각 조각을 독립적으로 캡처/재생하여 동적 형상을 지원한다.

이 글에서는 python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py를 중심으로 Piecewise CUDA Graph의 설계를 분석한다.

일반 CUDA Graph vs Piecewise CUDA Graph

두 접근의 핵심 차이를 비교한다.

일반 CUDA Graph (Decode):
┌───────────────────────────────────────┐
│  전체 모델이 하나의 그래프              │
│  Input(고정) → Layer1 → ... → Output  │
│  배치 크기별 1개 그래프                 │
└───────────────────────────────────────┘

Piecewise CUDA Graph (Prefill):
┌──────────┐  ┌──────────┐  ┌──────────┐
│ Piece 1  │  │ Piece 2  │  │ Piece 3  │
│ (Layers  │  │ (MoE     │  │ (Layers  │
│  0~11)   │──│ Dispatch)│──│ 12~23)   │
└──────────┘  └──────────┘  └──────────┘
  ↑ 각 조각이 독립 CUDA Graph
  ↑ 조각 사이에서 동적 연산 가능

PiecewiseCudaGraphRunner 초기화

초기화 과정은 크게 4단계이다.

class PiecewiseCudaGraphRunner:
    def __init__(self, model_runner: ModelRunner):
        self.model_runner = model_runner
        self.graphs = {}
        self.output_buffers = {}

        # 1. CompilationConfig 설정
        self.compile_config = CompilationConfig(
            self.model_runner.server_args.piecewise_cuda_graph_tokens,
            self.model_runner.server_args.piecewise_cuda_graph_compiler,
            self.model_runner.server_args.enable_torch_compile_debug_mode,
        )

        # 2. MoE 분할 지점 등록
        if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
            self.compile_config.add_split_op(
                "sglang.moe_forward_piecewise_cuda_graph_impl"
            )

        # 3. 캡처할 토큰 수 결정
        self.capture_num_tokens = self.compile_config.get_capture_sizes()

CompilationConfig의 split_ops가 분할 지점을 정의한다. DeepEP나 Mooncake 같은 All-to-All 통신이 필요한 MoE 레이어에서 그래프를 분할한다.

PrefillInputBuffers

Prefill용 입력 버퍼는 Decode와 다르게 가변 길이 입력을 지원해야 한다.

@dataclass
class PrefillInputBuffers(ForwardInputBuffers):
    input_ids: torch.Tensor
    out_cache_loc: torch.Tensor
    out_cache_loc_swa: Optional[torch.Tensor]
    mamba_track_indices: Optional[torch.Tensor]
    mamba_track_mask: Optional[torch.Tensor]
    mamba_track_seqlens: Optional[torch.Tensor]
    positions: torch.Tensor
    input_embeds: Optional[torch.Tensor]
    mrope_positions: Optional[torch.Tensor]

최대 토큰 수로 버퍼를 할당하고, 실제 사용 시 슬라이싱한다.

with torch.device(self.device):
    input_ids = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
    out_cache_loc = torch.zeros((self.max_num_tokens,), dtype=self._cache_loc_dtype())
    positions = torch.zeros((self.max_num_tokens,), dtype=torch.int64)

self.buffers = PrefillInputBuffers(
    input_ids=input_ids,
    out_cache_loc=out_cache_loc,
    positions=positions,
    ...
)

컴파일과 캡처

Piecewise CUDA Graph의 핵심은 torch.compile과 CUDA Graph 캡처를 결합하는 것이다. 초기화에서 3단계를 거친다.

1단계: Warmup - 첫 번째 토큰 수로 JIT 커널을 예열한다.

with enable_piecewise_cuda_graph():
    language_model = getattr(
        self.model_runner.model, "language_model", self.model_runner.model
    )
    with patch_model(
        language_model.model, self.compile_config.compiler
    ) as patched_model:
        self.warmup_compile(num_tokens=self.capture_num_tokens[0])

2단계: torch.compile 설치 - install_torch_compiled()로 모델에 컴파일 래퍼를 설치한다.

        install_torch_compiled(
            patched_model,
            fullgraph=True,
            dynamic_arg_dims=None,
            compile_config=self.compile_config,
            graph_pool=get_global_graph_memory_pool(),
        )

3단계: 모든 토큰 수에 대해 캡처 - 역순으로 컴파일하고, 마지막에 CUDA Graph를 캡처한다.

        with enable_piecewise_cuda_graph_compile():
            for num_tokens in reversed(self.capture_num_tokens):
                self.warmup_compile(num_tokens=num_tokens)

        # 최종 캡처
        self.capture()

warmup_compile

각 토큰 수에 대해 더미 ForwardBatch를 생성하고 포워드 패스를 실행한다.

def warmup_compile(self, num_tokens: int):
    buffers = self.buffers
    input_ids = buffers.input_ids[:num_tokens]
    positions = buffers.positions[:num_tokens]
    out_cache_loc = buffers.out_cache_loc[:num_tokens]

    forward_batch = ForwardBatch(
        forward_mode=ForwardMode.EXTEND,
        batch_size=1,
        input_ids=input_ids,
        seq_lens=torch.tensor([num_tokens], device=self.device),
        extend_num_tokens=num_tokens,
        ...
    )

    self.model_runner.attn_backend.init_forward_metadata(forward_batch)
    ...

Eager vs Inductor

컴파일러는 두 가지 옵션을 지원한다.

assert self.model_runner.server_args.piecewise_cuda_graph_compiler in [
    "eager",
    "inductor",
], "By now, only eager and inductor are supported"
컴파일러 특징 사용 시기
eager 컴파일 없이 CUDA Graph만 캡처 빠른 초기화가 중요할 때
inductor torch.compile + CUDA Graph 최대 성능이 필요할 때

patch_model()에서 컴파일러에 따라 MultiPlatformOp을 토글한다.

@contextmanager
def patch_model(model, compiler):
    try:
        if compiler != "eager":
            _to_torch(model, reverse=False, num_tokens=16)
        yield model
    finally:
        _to_torch(model, reverse=True, num_tokens=16)

메모리 풀 공유

일반 CudaGraphRunner와 동일한 글로벌 메모리 풀을 공유한다.

if get_global_graph_memory_pool() is None:
    set_global_graph_memory_pool(self.device_module.graph_pool_handle())
set_graph_pool_id(get_global_graph_memory_pool())

이를 통해 Decode용 CUDA Graph와 Prefill용 Piecewise Graph가 메모리를 효율적으로 공유한다.

설계 근거: 분할 지점의 선택

왜 MoE All-to-All 통신에서 분할하는가? CUDA Graph는 NCCL 통신 패턴이 고정되어야 한다. DeepEP/Mooncake의 동적 라우팅은 이 조건을 만족하지 않으므로, 해당 지점에서 그래프를 분할해야 한다.

┌──────────────────┐     ┌──────────────────┐
│  CUDA Graph 1    │     │  CUDA Graph 2    │
│  Attention +     │     │  Attention +     │
│  FFN Layers      │     │  FFN Layers      │
└────────┬─────────┘     └──────────────────┘
         │                        ↑
         ▼                        │
┌──────────────────┐              │
│  MoE Dispatch    │──────────────┘
│  (동적 라우팅,    │
│   CUDA Graph 외부)│
└──────────────────┘

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글