본문으로 건너뛰기

[SGLang] Outlines: FSM 기반 제약 생성과 Jump-Forward 최적화

들어가며

Outlines는 SGLang이 최초로 통합한 Constrained Decoding 백엔드다. 정규식을 Finite State Machine(FSM)으로 변환하고, JSON Schema도 정규식으로 변환한 뒤 FSM을 구축하는 방식을 취한다. 특히 SGLang에서는 Compressed FSM(Jump-Forward Decoding) 이라는 최적화를 추가하여, 다음 토큰이 확정적일 때 여러 토큰을 한 번에 건너뛸 수 있다.

소스 파일은 python/sglang/srt/constrained/outlines_backend.pyoutlines_jump_forward.py 두 개로 구성된다.

구조도

┌───────────────────────────────────────────────┐
│           OutlinesGrammarBackend              │
│  ┌────────────────────┐                       │
│  │TransformerTokenizer│  (outlines 어댑터)     │
│  └─────────┬──────────┘                       │
│            │                                  │
│  dispatch_json()                              │
│    └─► build_regex_from_schema()              │
│          └─► _compile_regex()                 │
│                                               │
│  dispatch_regex()                             │
│    └─► _compile_regex()                       │
│          └─► RegexGuide (FSM)                 │
│                │                              │
│                ▼                              │
│  ┌────────────────────────┐                   │
│  │   OutlinesGrammar      │                   │
│  │   state: int (FSM)     │                   │
│  │   guide: RegexGuide    │                   │
│  └──────────┬─────────────┘                   │
│             │                                 │
│             ▼                                 │
│  ┌────────────────────────┐                   │
│  │ OutlinesJumpForwardMap │  (optional)       │
│  │ state → JumpEdge       │                   │
│  └────────────────────────┘                   │
└───────────────────────────────────────────────┘

핵심 코드 분석

1. JSON Schema를 Regex로 변환

class OutlinesGrammarBackend(BaseGrammarBackend):
    def dispatch_json(self, key_string: str):
        try:
            regex = build_regex_from_object(
                key_string,
                whitespace_pattern=self.whitespace_pattern,
            )
        except (NotImplementedError, json.decoder.JSONDecodeError, ValueError) as e:
            return InvalidGrammarObject(str(e))
        return self._compile_regex(regex)

Outlines의 핵심 전략: JSON Schema를 먼저 정규식으로 변환한 뒤, 정규식을 FSM으로 컴파일한다. 이 2단계 변환 파이프라인은 구현이 단순하지만, 복잡한 JSON Schema(재귀 구조, anyOf 등)에서 정규식이 폭발적으로 커질 수 있다.

def build_regex_from_object(object, whitespace_pattern=None):
    if isinstance(object, type(BaseModel)):
        schema = json.dumps(object.model_json_schema())
    elif isinstance(object, Dict):
        schema = json.dumps(object)
    else:
        schema = object
    return build_regex_from_schema(schema, whitespace_pattern)

Pydantic BaseModel, dict, 문자열 등 다양한 입력 형식을 지원한다.

2. Regex에서 FSM 구축

def _compile_regex(self, regex: str) -> BaseGrammarObject:
    try:
        if hasattr(RegexGuide, "from_regex"):
            guide = RegexGuide.from_regex(regex, self.outlines_tokenizer)
        else:
            guide = RegexGuide(regex, self.outlines_tokenizer)
    except interegular.patterns.InvalidSyntax as e:
        return InvalidGrammarObject(str(e))
    return OutlinesGrammar(guide, jump_forward_map=None)

RegexGuide는 정규식과 토크나이저를 받아 FSM을 구축한다. FSM의 각 상태에서 허용되는 토큰 목록을 미리 계산하므로, 런타임에는 상태 전이와 토큰 마스킹만 수행한다. Outlines 버전 호환성을 위해 from_regex와 직접 생성자 두 방식을 모두 지원한다.

3. Bool 텐서 기반 마스킹

class OutlinesGrammar(BaseGrammarObject):
    def accept_token(self, token: int):
        self.state = self.guide.get_next_state(self.state, token)

    def fill_vocab_mask(self, vocab_mask, idx):
        tokens = torch.tensor(
            self.guide.get_next_instruction(self.state).tokens, dtype=torch.int64
        ).to(vocab_mask.device, non_blocking=True)
        vocab_mask = vocab_mask[idx]
        vocab_mask.fill_(1)
        vocab_mask.scatter_(0, tokens, torch.zeros_like(tokens, dtype=torch.bool))

    @staticmethod
    def apply_vocab_mask(logits, vocab_mask):
        logits.masked_fill_(vocab_mask, float("-inf"))

XGrammar의 비트마스크와 달리, Outlines는 Bool 텐서를 사용한다. fill_(1)scatter_(허용 토큰, 0) 패턴으로 허용 토큰 위치만 False로 남기고, 나머지를 True로 설정한다. 이후 masked_fill_True 위치의 logits를 -inf로 만든다.

4. Jump-Forward Decoding (Compressed FSM)

Jump-Forward는 SGLang이 제안한 핵심 최적화다. FSM의 특정 상태에서 나가는 간선이 하나뿐이면, 다음 토큰이 확정적이므로 LLM 추론 없이 바로 건너뛸 수 있다.

# outlines_jump_forward.py
@dataclasses.dataclass
class JumpEdge:
    symbol: str = None
    symbol_next_state: int = None
    byte: int = None
    byte_next_state: int = None

JumpEdge는 확정적 전이를 표현한다. symbol은 문자 수준, byte는 바이트 수준 전이를 저장한다. UTF-8 멀티바이트 문자 처리를 위해 두 수준을 모두 관리한다.

def init_state_to_jump_forward(regex_string):
    byte_fsm = make_byte_level_fsm(regex_pattern.to_fsm().reduce(), keep_utf8=True)
    regex_fsm, _ = make_deterministic_fsm(byte_fsm)
    fsm_info: FSMInfo = regex_fsm.fsm_info

    outgoings_ct = defaultdict(int)
    for s in fsm_info.finals:
        outgoings_ct[s] = 1  # 종결 상태는 이미 나가는 간선 1개

    state_to_jump_forward = {}
    for (state, id_), next_state in transitions.items():
        if id_ == fsm_info.alphabet_anything_value:
            continue  # 와일드카드는 점프 불가

        symbols = id_to_symbol[id_]
        for c in symbols:
            if len(c) > 1:
                continue
            outgoings_ct[state] += 1
            if outgoings_ct[state] > 1:
                if state in state_to_jump_forward:
                    del state_to_jump_forward[state]
                break

            state_to_jump_forward[state] = JumpEdge(
                symbol=c, symbol_next_state=next_state,
            )

FSM의 각 상태를 순회하며 나가는 간선 수를 센다. 간선이 정확히 1개인 상태만 state_to_jump_forward에 등록한다. 와일드카드(anything_value)는 임의 입력을 허용하므로 점프 대상에서 제외한다.

5. 연쇄 점프

class OutlinesJumpForwardMap:
    def jump_forward_symbol(self, state):
        jump_forward_str = ""
        next_state = state
        while state in self.state_to_jump_forward:
            e = self.state_to_jump_forward[state]
            if e.symbol is None:
                break
            jump_forward_str += e.symbol
            next_state = e.symbol_next_state
            state = next_state
        return jump_forward_str, next_state

연속된 확정 상태를 체인으로 따라가며 한 번에 여러 문자를 건너뛴다. 예를 들어 "The google's DNS server address is " 같은 고정 prefix가 있으면 전체 문자열을 한 번에 점프한다.

6. 바이트 수준 Jump-Forward

def jump_forward_byte(self, state):
    jump_forward_bytes = []
    while state in self.state_to_jump_forward:
        e = self.state_to_jump_forward[state]
        jump_forward_bytes.append((e.byte, e.byte_next_state))
        next_state = e.byte_next_state
        state = next_state
    return jump_forward_bytes

UTF-8 문자가 여러 바이트로 구성될 때, 바이트 수준에서도 점프를 수행한다. OutlinesGrammar.try_jump_forward()에서 continuation byte(0x80~0xBF)를 먼저 처리한 뒤 기호 수준 점프를 실행한다.

XGrammar와의 비교

항목 Outlines XGrammar
문법 모델 FSM (정규 문법만) PDA (Context-Free Grammar)
JSON 처리 Schema → Regex → FSM Schema → 직접 PDA 컴파일
EBNF 미지원 지원
마스크 형식 Bool 텐서 (32x 메모리) 비트마스크
Jump-Forward Compressed FSM (바이트/기호 수준) find_jump_forward_string()
Rollback 상태 변수 직접 설정 matcher.rollback(k)

Outlines는 정규 문법 수준의 제약만 표현할 수 있지만, FSM이 단순하여 디버깅이 쉽고 Jump-Forward의 구현이 직관적이다. 복잡한 재귀적 JSON Schema가 필요하면 XGrammar나 LLGuidance가 더 적합하다.

관련 포스트

참고

댓글

관련 포스트

SGLang 의 다른글