본문으로 건너뛰기

[SGLang] Grammar Manager: 구조화된 출력 생성의 통합 관리

들어가며

LLM이 JSON, SQL, 코드 등 특정 형식의 출력을 생성해야 할 때, 자유 생성(free-form generation)은 형식 오류를 일으킨다. SGLang은 Constrained Decoding 레이어를 통해 매 토큰 생성 시 허용 가능한 토큰만 남기는 방식으로 이 문제를 해결한다. 이 레이어의 최상위에 위치하는 것이 GrammarManager다.

GrammarManager는 요청에 포함된 제약 조건(JSON Schema, Regex, EBNF, Structural Tag)을 감지하고, 적절한 백엔드에 문법 컴파일을 위임하며, 비동기 컴파일 완료를 관리하는 오케스트레이터 역할을 한다.

구조도

┌─────────────────────────────────────────────┐
│                GrammarManager               │
│  ┌───────────┐  ┌────────────────────────┐  │
│  │ grammar   │  │ grammar_backend        │  │
│  │ _queue    │  │ (BaseGrammarBackend)   │  │
│  │ [Req...]  │  │                        │  │
│  └─────┬─────┘  └────────┬───────────────┘  │
│        │                 │                  │
│        │    ┌────────────┼────────────┐     │
│        │    │            │            │     │
│        ▼    ▼            ▼            ▼     │
│     ┌──────────┐  ┌──────────┐  ┌────────┐ │
│     │XGrammar  │  │Outlines  │  │LLGuid. │ │
│     │Backend   │  │Backend   │  │Backend │ │
│     └──────────┘  └──────────┘  └────────┘ │
│                                             │
│  ┌─────────────────────────────────────┐    │
│  │ ReasonerGrammarBackend (wrapper)    │    │
│  │ - think_end_id로 추론/응답 분리     │    │
│  └─────────────────────────────────────┘    │
└─────────────────────────────────────────────┘

핵심 코드 분석

1. 백엔드 생성과 레지스트리

GrammarManager는 초기화 시 create_grammar_backend()를 호출하여 서버 설정에 맞는 백엔드를 생성한다.

# base_grammar_backend.py
def create_grammar_backend(
    server_args, tokenizer, vocab_size, eos_token_ids, think_end_id=None
) -> Optional[BaseGrammarBackend]:
    name = server_args.grammar_backend

    if name in GRAMMAR_BACKEND_REGISTRY:
        return GRAMMAR_BACKEND_REGISTRY[name](...)

    if name == "xgrammar":
        grammar_backend = XGrammarGrammarBackend(tokenizer, vocab_size=vocab_size, ...)
    elif name == "outlines":
        grammar_backend = OutlinesGrammarBackend(tokenizer, ...)
    elif name == "llguidance":
        grammar_backend = GuidanceBackend(tokenizer=tokenizer, ...)
    elif name == "none":
        return None

GRAMMAR_BACKEND_REGISTRY를 통해 커스텀 백엔드를 등록할 수 있으며, 기본 백엔드 3종(xgrammar, outlines, llguidance) 중 하나를 선택한다. reasoning_parser가 설정되어 있으면 ReasonerGrammarBackend로 한 번 더 래핑한다.

2. 요청의 제약 조건 감지와 큐잉

# grammar_manager.py
def process_req_with_grammar(self, req: Req) -> bool:
    if (
        req.sampling_params.json_schema is not None
        or req.sampling_params.regex is not None
        or req.sampling_params.ebnf is not None
        or req.sampling_params.structural_tag is not None
    ):
        if req.sampling_params.json_schema is not None:
            key = ("json", req.sampling_params.json_schema)
        elif req.sampling_params.regex is not None:
            key = ("regex", req.sampling_params.regex)
        # ...
        value, cache_hit = self.grammar_backend.get_cached_or_future_value(
            key, req.require_reasoning
        )
        req.grammar = value

요청의 sampling_params에서 제약 조건을 감지하고, (타입, 값) 튜플을 키로 만든다. get_cached_or_future_value()는 캐시 히트 시 복사본을 즉시 반환하고, 미스 시 ThreadPoolExecutor에 컴파일을 제출한 뒤 Future를 반환한다.

3. 캐시와 비동기 컴파일

# base_grammar_backend.py
def get_cached_or_future_value(self, key, require_reasoning):
    value = self.cache.get(key)
    if value:
        copied_value = value.copy()
        copied_value.maybe_init_reasoning(require_reasoning)
        return copied_value, True
    value = self.executor.submit(self._init_value_dispatch, key, require_reasoning)
    return value, False

캐시 히트 시 .copy()로 독립적인 상태의 Grammar 객체를 반환한다. 이는 각 요청이 자신만의 FSM 상태를 유지해야 하기 때문이다. 캐시 미스 시 _init_value_dispatch()가 백그라운드 스레드에서 실행되며, 제약 타입에 따라 dispatch_json(), dispatch_regex(), dispatch_ebnf(), dispatch_structural_tag() 중 하나를 호출한다.

4. 비동기 폴링과 분산 동기화

# grammar_manager.py
def get_ready_grammar_requests(self) -> List[Req]:
    start_time = time.perf_counter()
    while time.perf_counter() - start_time < self.SGLANG_GRAMMAR_POLL_INTERVAL:
        for i, req in enumerate(self.grammar_queue):
            if isinstance(req.grammar, futures.Future) and req.grammar.done():
                ready_req_idxs.add(i)
        time.sleep(self.SGLANG_GRAMMAR_POLL_INTERVAL / 10)

스케줄러는 GPU forward pass 사이에 get_ready_grammar_requests()를 호출하여 컴파일이 완료된 요청을 수거한다. SGLANG_GRAMMAR_POLL_INTERVAL 동안 폴링하며, SGLANG_GRAMMAR_MAX_POLL_ITERATIONS 횟수를 초과하면 타임아웃으로 처리한다.

5. 다중 랭크 동기화

if self.grammar_sync_size == 1:
    synced_ready_req_idxs = ready_req_idxs
else:
    all_gather_output = [None] * self.grammar_sync_size
    torch.distributed.all_gather_object(
        all_gather_output,
        (ready_req_idxs, failed_req_idxs),
        group=self.grammar_sync_group,
    )
    synced_ready_req_idxs = set.intersection(*[x[0] for x in all_gather_output])
    synced_failed_req_idxs = set.union(*[x[1] for x in all_gather_output])

Data Parallel + Tensor Parallel 환경에서 모든 랭크가 동일한 요청을 처리하려면 동기화가 필요하다. ready 집합은 교집합(모든 랭크에서 준비 완료된 것만), failed 집합은 합집합(어느 랭크에서든 실패하면 전체 실패)으로 처리한다.

백엔드 비교표

항목 XGrammar Outlines LLGuidance
개발 SGLang/MLC 커뮤니티 Outlines 프로젝트 Microsoft
JSON Schema GrammarCompiler.compile_json_schema() build_regex_from_schema() → Regex LLMatcher.grammar_from_json_schema()
Regex compile_regex() RegexGuide (FSM 구축) grammar_from("regex", ...)
EBNF compile_grammar() 미지원 grammar_from("ebnf", ...)
Structural Tag compile_structural_tag() 미지원 StructTag.to_grammar()
마스크 방식 비트마스크 (Triton/CUDA) Bool 텐서 + masked_fill_ 비트마스크 (llguidance.torch)
Jump-Forward find_jump_forward_string() Compressed FSM compute_ff_tokens()
Rollback matcher.rollback(k) 상태 직접 설정 ll_matcher.rollback(k)
기본값 SGLang 기본 백엔드 레거시 설치 필요

설계 근거

왜 비동기 컴파일인가? JSON Schema를 문법으로 컴파일하는 데 수십~수백 ms가 소요된다. GPU forward pass를 블로킹하지 않기 위해 ThreadPoolExecutor에서 병렬로 컴파일하고, 스케줄러는 다른 요청을 먼저 처리한다.

왜 캐시에서 copy()인가? Grammar 객체는 FSM 상태를 내부에 유지한다. 동일한 JSON Schema를 사용하는 두 요청이 서로 다른 토큰 시퀀스를 생성하므로, 캐시된 컴파일 결과를 공유하되 상태는 독립적이어야 한다.

왜 교집합/합집합 동기화인가? 준비 완료는 모든 랭크에서 확인되어야 배치에 포함할 수 있다(교집합). 반면 실패는 어느 한 랭크에서라도 발생하면 해당 요청 전체를 중단해야 한다(합집합).

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글