[vLLM] RoPE 변형: 15+ 로타리 위치 인코딩

들어가며

Rotary Position Embedding(RoPE)은 현대 LLM의 위치 인코딩 표준이 되었다. GPT-NeoX, LLaMA, Mistral, Qwen, DeepSeek 등 거의 모든 주요 모델이 RoPE를 채택한다. 그런데 "RoPE"라고 한 마디로 부르기엔, 실제로는 십수 가지 변형이 존재한다. 컨텍스트 길이 확장, 다중 해상도 지원, 비전-언어 모델 대응 등 각각의 요구사항에 맞춘 변형이 vLLM에 모두 구현되어 있다.

• 논문: RoFormer: Enhanced Transformer with Rotary Position Embedding
• 공식 문서: https://docs.vllm.ai

핵심 구조/코드 분석

전체 변형 목록

vllm/model_executor/layers/rotary_embedding/ 디렉토리에 각 변형이 독립 파일로 구현되어 있다:

파일	변형	대표 모델
base.py	기본 RoPE	GPT-NeoX, Pythia
linear_scaling_rope.py	Linear Scaling	초기 컨텍스트 확장
dynamic_ntk_scaling_rope.py	Dynamic NTK Scaling	CodeLlama
dynamic_ntk_alpha_rope.py	Dynamic NTK Alpha	-
ntk_scaling_rope.py	NTK Scaling	-
yarn_scaling_rope.py	YaRN Scaling	Yarn 모델
llama3_rope.py	Llama3 RoPE	Llama 3/3.1
deepseek_scaling_rope.py	DeepSeek Scaling	DeepSeek-V2/V3
mrope.py	Multi-dimensional RoPE	Qwen2-VL
mrope_interleaved.py	MRoPE Interleaved	-
dual_chunk_rope.py	Dual Chunk RoPE	-
phi3_long_rope_scaled_rope.py	Phi3 Long RoPE	Phi-3
gemma4_rope.py	Gemma4 RoPE	Gemma 4
llama4_vision_rope.py	Llama4 Vision RoPE	Llama 4 Vision
fope.py	Fourier RoPE	-
xdrope.py	XD RoPE	-
telechat3_scaling_rope.py	TeleChat3 Scaling	TeleChat3

기본 구현: RotaryEmbeddingBase

모든 변형의 부모 클래스이다:

@CustomOp.register("rotary_embedding")
class RotaryEmbeddingBase(CustomOp):
    def __init__(
        self,
        head_size: int,
        rotary_dim: int,
        max_position_embeddings: int,
        base: float,
        is_neox_style: bool,
        dtype: torch.dtype,
    ) -> None:
        self.head_size = head_size
        self.rotary_dim = rotary_dim
        self.base = base
        self.is_neox_style = is_neox_style

    def _compute_inv_freq(self, base: float) -> torch.Tensor:
        inv_freq = 1.0 / (
            base ** (
                torch.arange(0, self.rotary_dim, 2, dtype=torch.float)
                / self.rotary_dim
            )
        )
        return inv_freq

    def _compute_cos_sin_cache(self) -> torch.Tensor:
        inv_freq = self._compute_inv_freq(self.base)
        t = torch.arange(self.max_position_embeddings, dtype=torch.float)
        freqs = torch.einsum("i,j -> ij", t, inv_freq)
        cos = freqs.cos()
        sin = freqs.sin()
        cache = torch.cat((cos, sin), dim=-1)
        return cache

핵심은 역주파수 벡터(inv_freq)를 미리 계산하고, 위치별 cos/sin 캐시를 만들어두는 것이다. is_neox_style은 회전을 적용하는 방식(GPT-NeoX 스타일 vs GPT-J 스타일)을 결정한다.

get_rope: 팩토리 함수

모든 RoPE 변형을 하나의 인터페이스로 생성하는 팩토리가 있다:

def get_rope(
    head_size: int,
    max_position: int,
    is_neox_style: bool = True,
    rope_parameters: dict[str, Any] | None = None,
    dtype: torch.dtype | None = None,
    dual_chunk_attention_config: dict[str, Any] | None = None,
) -> RotaryEmbedding:
    rope_parameters = rope_parameters or {}
    base = rope_parameters.get("rope_theta", 10000)
    scaling_type = rope_parameters.get("rope_type", "default")
    partial_rotary_factor = rope_parameters.get("partial_rotary_factor", 1.0)
    rotary_dim = int(head_size * partial_rotary_factor)

_ROPE_DICT로 동일한 설정의 RoPE 인스턴스를 캐싱하여, 같은 파라미터로 중복 생성을 방지한다. scaling_type에 따라 적절한 하위 클래스를 인스턴스화한다.

하드웨어 가속

RoPE는 CustomOp으로 등록되어 다양한 백엔드를 활용한다:

self.use_aiter = (
    self.enabled() and rocm_aiter_ops.is_triton_rotary_embed_enabled()
)

ROCm에서는 AITER의 Triton 기반 RoPE 커널을, NVIDIA에서는 FlashInfer의 RoPE 커널을 선택적으로 사용한다. 이 분기는 각 플랫폼에서 최적의 성능을 보장한다.

왜 이 설계인가

1. cos/sin 캐시 프리컴퓨팅: RoPE 연산은 위치 인덱스에 대한 삼각함수 계산을 요구한다. 이를 매 토큰마다 계산하면 낭비이므로, 최대 위치까지의 값을 미리 계산하여 캐시한다.

2. partial_rotary_factor: 일부 모델(Phi 계열)은 전체 head_size가 아닌 일부 차원에만 RoPE를 적용한다. 나머지 차원은 위치 무관(position-invariant)으로 남겨 다른 정보를 인코딩하는 데 사용한다.

3. 캐싱 딕셔너리: 동일한 설정의 RoPE가 여러 레이어에서 사용될 때, _ROPE_DICT를 통해 하나의 인스턴스만 생성한다. 이는 cos/sin 캐시의 메모리를 공유하는 효과가 있다.

4. 변형별 독립 파일: 각 RoPE 변형은 수학적으로 다른 주파수 스케일링을 적용한다. 독립 파일로 분리하면 새 모델이 새로운 RoPE 변형을 요구할 때 기존 코드 수정 없이 확장할 수 있다.

논문 핵심 내용

RoFormer: Enhanced Transformer with Rotary Position Embedding (2104.09864) 논문은 회전 행렬을 이용한 위치 인코딩 기법을 제안했다.

핵심 아이디어: RoPE는 절대 위치를 회전 행렬로 인코딩하면서, 동시에 셀프 어텐션 내적에서 명시적인 상대 위치 의존성을 자연스럽게 만들어낸다. 이전의 절대 위치 인코딩(sinusoidal)이나 상대 위치 인코딩(T5 bias, ALiBi)과 달리, RoPE는 절대 + 상대 위치 정보를 하나의 메커니즘으로 통합했다.

RoPE의 핵심 장점

특성	설명
시퀀스 길이 유연성	학습 시 길이 제한 없이 임의 길이로 확장 가능
거리 기반 감쇠	토큰 간 상대 거리가 멀수록 의존도가 자연스럽게 감소
선형 어텐션 호환	상대 위치를 가진 선형 셀프 어텐션 구현 가능
추가 파라미터 없음	학습 가능한 파라미터가 0개

논문의 실험 결과에서 RoFormer는 장문 텍스트 분류 벤치마크에서 기존 대안(절대 위치 인코딩, 상대 위치 인코딩)을 일관되게 능가했다. 특히 시퀀스 길이가 길어질수록 성능 격차가 벌어진다는 점이 중요하다.

RoPE가 사실상의 표준이 된 이유는 수학적 우아함과 실용성의 조합이다. cos/sin 캐시만 미리 계산하면 추론 시 추가 파라미터 없이 동작하고, 회전이라는 직관적인 연산이라 하드웨어 가속에도 유리하다. GPT-NeoX에서 처음 채택된 이후 LLaMA, Mistral, Qwen, DeepSeek 등 거의 모든 주요 모델이 RoPE를 기본으로 사용하게 됐다.

마무리

vLLM은 15가지 이상의 RoPE 변형을 하나의 통합된 인터페이스 아래에서 관리한다. 기본 RoPE에서 YaRN, Llama3, 멀티모달 MRoPE까지, 각 모델의 위치 인코딩 요구사항을 그대로 반영하면서도 cos/sin 캐시 공유와 하드웨어 가속을 통해 효율적으로 실행한다.