[vLLM] Structured Output: JSON/regex/문법 제약 생성
들어가며
LLM의 출력을 프로그래밍 방식으로 소비하려면 구조화된 형식이 필요하다. "JSON으로 답변해줘"라는 프롬프트만으로는 유효한 JSON이 나온다는 보장이 없다. Structured Output은 디코딩 과정에서 문법 규칙을 강제하여, 생성되는 모든 토큰이 지정된 스키마를 만족하도록 보장한다. vLLM v1에서는 이 기능이 vllm/v1/structured_output/에 구현되어 있으며, xgrammar를 기본 백엔드로 사용한다.
- 논문: XGrammar: Flexible and Efficient Structured Generation Engine for Large Language Models
- 공식 문서: https://docs.vllm.ai
공식 문서
vLLM 공식 문서: Structured Outputs
핵심 구조/코드 분석
백엔드 아키텍처
vllm/v1/structured_output/ 디렉토리에 5개의 백엔드가 구현되어 있다:
| 파일 | 백엔드 | 특징 |
|---|---|---|
backend_xgrammar.py |
XGrammar | 기본 백엔드, C++ 기반 고성능 |
backend_outlines.py |
Outlines | 순수 Python, 폭넓은 문법 지원 |
backend_guidance.py |
Guidance | Microsoft의 제약 생성 엔진 |
backend_lm_format_enforcer.py |
LM Format Enforcer | - |
backend_types.py |
공통 타입 정의 | 추상 인터페이스 |
StructuredOutputOptions: 지원 형식
class StructuredOutputOptions(enum.Enum):
JSON = enum.auto()
JSON_OBJECT = enum.auto()
REGEX = enum.auto()
GRAMMAR = enum.auto()
CHOICE = enum.auto()
STRUCTURAL_TAG = enum.auto()
6가지 제약 유형을 지원한다. JSON은 JSON 스키마를, REGEX는 정규식 패턴을, GRAMMAR는 EBNF/Lark 문법을, CHOICE는 선택지 목록을 강제한다. STRUCTURAL_TAG는 태그 기반 구조를 지원하는 고급 기능이다.
StructuredOutputGrammar: 추상 인터페이스
모든 백엔드의 문법 객체가 구현해야 할 인터페이스이다:
class StructuredOutputGrammar(ABC):
@abstractmethod
def accept_tokens(self, request_id: str, tokens: list[int]) -> bool:
"""토큰이 문법에 의해 수락되는지 확인한다."""
@abstractmethod
def validate_tokens(self, tokens: list[int]) -> list[int]:
"""FSM을 전진시키지 않고 토큰을 검증한다."""
@abstractmethod
def rollback(self, num_tokens: int) -> None:
"""지정된 수만큼 토큰을 롤백한다."""
@abstractmethod
def fill_bitmask(self, bitmask: torch.Tensor, batch_index: int) -> None:
"""배치 인덱스에 대한 비트마스크를 채운다."""
fill_bitmask가 핵심이다. 각 디코딩 스텝에서 현재 문법 상태에서 허용되는 토큰의 비트마스크를 생성하고, 이를 로짓에 적용하여 불허 토큰의 확률을 -inf로 만든다. rollback은 speculative decoding에서 거부된 토큰을 되돌리기 위해 사용된다.
XGrammar 백엔드: 기본 구현
@dataclass
class XgrammarBackend(StructuredOutputBackend):
def __post_init__(self):
tokenizer_info = xgr.TokenizerInfo.from_huggingface(
self.tokenizer, vocab_size=self.vocab_size,
)
self.compiler = xgr.GrammarCompiler(
tokenizer_info,
max_threads=8,
cache_enabled=True,
cache_limit_bytes=vllm.envs.VLLM_XGRAMMAR_CACHE_MB * 1024 * 1024,
)
def compile_grammar(self, request_type, grammar_spec):
if request_type == StructuredOutputOptions.JSON:
ctx = self.compiler.compile_json_schema(
grammar_spec,
any_whitespace=not self.disable_any_whitespace
)
elif request_type == StructuredOutputOptions.REGEX:
ctx = self.compiler.compile_regex(grammar_spec)
elif request_type == StructuredOutputOptions.GRAMMAR:
ctx = self.compiler.compile_grammar(grammar_spec)
...
return XgrammarGrammar(
matcher=xgr.GrammarMatcher(
ctx, max_rollback_tokens=self.num_speculative_tokens,
),
vocab_size=self.vocab_size, ctx=ctx,
)
GrammarCompiler가 JSON 스키마/정규식/문법을 유한 상태 머신(FSM)으로 컴파일하고, GrammarMatcher가 런타임에 FSM을 추적하며 비트마스크를 생성한다. 캐시가 활성화되어 동일한 스키마 요청을 반복 컴파일하지 않는다.
Speculative Decoding 호환
self.num_speculative_tokens = 0
if self.vllm_config.speculative_config is not None:
self.num_speculative_tokens = (
self.vllm_config.speculative_config.num_speculative_tokens
)
max_rollback_tokens을 speculative token 수로 설정하여, 투기적 디코딩에서 거부된 토큰을 올바르게 롤백할 수 있다.
왜 이 설계인가
1. 비트마스크 기반 제약: 로짓에 대한 마스킹은 추론 결과의 품질에 영향을 주지 않으면서 100% 유효한 출력을 보장한다. 재시도 기반 접근보다 훨씬 효율적이다.
2. 멀티 백엔드: xgrammar는 빠르지만 모든 문법을 지원하지 않을 수 있다. outlines는 느리지만 더 유연하다. 사용자가 요구사항에 맞는 백엔드를 선택할 수 있게 하는 것이 실용적이다.
3. 컴파일-실행 분리: 스키마를 FSM으로 컴파일하는 과정은 비용이 크므로 캐싱이 중요하다. cache_enabled=True와 VLLM_XGRAMMAR_CACHE_MB를 통해 컴파일 결과를 재사용한다.
4. Mistral 토크나이저 특수 처리: Mistral의 Tekken 토크나이저는 표준 BPE와 다른 인코딩을 사용하므로, VocabType.RAW로 별도 처리한다. 이런 세밀한 호환성 처리가 프로덕션 환경에서의 안정성을 보장한다.
논문 핵심 내용
XGrammar: Flexible and Efficient Structured Generation Engine (2411.15100) 논문은 LLM 구조화 출력 생성을 위한 고성능 문법 엔진을 제안했다.
핵심 아이디어: 토큰을 컨텍스트 독립 토큰(미리 체크 가능)과 컨텍스트 의존 토큰(런타임에 해석 필요)으로 분할하는 어휘 파티셔닝 기법이 핵심이다. 문법 변환으로 컨텍스트를 확장하고 컨텍스트 독립 토큰 수를 줄이며, LLM 추론 엔진과 공동 설계하여 문법 연산을 GPU 실행과 오버랩시킨다.
마스크 생성 지연시간
| 문법 유형 | XGrammar | 기존 방법 대비 |
|---|---|---|
| JSON Schema | < 40 us/토큰 | 최대 3x 빠름 |
| Context-free Grammar (JSON) | - | 100x 이상 빠름 |
| XML | < 200 us/토큰 | - |
| Python DSL | < 200 us/토큰 | - |
최적화 단계별 성능 개선 (Ablation)
| 최적화 단계 | 지연시간 | 누적 개선 |
|---|---|---|
| PDA Baseline | 65.776 ms | 1x |
| + Node merging | 38.280 ms | 1.7x |
| + Adaptive token mask cache | 0.154 ms | 427x |
| + Rule inlining | 0.035 ms | 1,879x |
| + Context Expansion | 0.018 ms | 3,654x |
엔드투엔드 서빙에서 H100 GPU 기준 최대 80x 속도 향상을 달성했다. Llama-3.1-8B에서 JSON Schema 생성 시 SGLang + XGrammar 조합은 배치 1에서 토큰당 6.2-6.3ms를 기록했는데, 최적화 없는 경우 44.2ms였다. 또한 제약 조건 적용으로 function calling 정확도가 62%에서 100%로, XML 코드 생성은 80%에서 100%로 향상됐다.
마무리
vLLM의 Structured Output은 비트마스크 기반 제약 디코딩을 통해 LLM 출력의 형식을 100% 보장한다. xgrammar의 고성능 FSM 컴파일러를 기본으로 하되, outlines/guidance 등 대안 백엔드도 제공하여 다양한 사용 시나리오를 커버한다.
관련 포스트
vLLM 의 다른글
- 이전글 [vLLM] RoPE 변형: 15+ 로타리 위치 인코딩
- 현재글 : [vLLM] Structured Output: JSON/regex/문법 제약 생성
- 다음글 [vLLM] Automatic Prefix Caching: 접두사 캐싱
댓글