Split View: Mamba vs Transformer 실무 비교: 논문에서 제품까지

"""
Mamba vs Transformer 실측 벤치마크 러너.

동일한 입력 세트에 대해 두 모델의 지연, 처리량, 메모리, 품질을
통제된 환경에서 측정한다.
"""
import time
import torch
import json
from dataclasses import dataclass, asdict
from pathlib import Path


@dataclass
class BenchmarkResult:
    model_name: str
    sequence_length: int
    batch_size: int
    # 성능 지표
    ttft_ms: float           # Time to First Token
    tps: float               # Tokens per Second (생성)
    total_latency_ms: float  # 전체 응답 지연
    # 자원 지표
    peak_vram_gb: float
    avg_vram_gb: float
    # 품질 지표 (별도 평가)
    quality_score: float = 0.0


class ModelBenchmarker:
    def __init__(self, model, tokenizer, model_name: str):
        self.model = model
        self.tokenizer = tokenizer
        self.model_name = model_name

    @torch.no_grad()
    def measure_generation(
        self,
        prompt: str,
        max_new_tokens: int = 256,
    ) -> BenchmarkResult:
        """단일 프롬프트에 대한 생성 성능을 측정한다."""
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt")
        input_ids = input_ids.to(self.model.device)
        seq_len = input_ids.shape[1]

        torch.cuda.reset_peak_memory_stats()

        # TTFT 측정: 첫 토큰 생성까지의 시간
        start = time.perf_counter()
        outputs = self.model.generate(
            input_ids,
            max_new_tokens=1,
            do_sample=False,
        )
        ttft = (time.perf_counter() - start) * 1000

        # 전체 생성 측정
        torch.cuda.synchronize()
        start = time.perf_counter()
        outputs = self.model.generate(
            input_ids,
            max_new_tokens=max_new_tokens,
            do_sample=False,
        )
        total_time = (time.perf_counter() - start) * 1000

        generated_tokens = outputs.shape[1] - seq_len
        tps = generated_tokens / (total_time / 1000) if total_time > 0 else 0

        peak_vram = torch.cuda.max_memory_allocated() / (1024**3)
        current_vram = torch.cuda.memory_allocated() / (1024**3)

        return BenchmarkResult(
            model_name=self.model_name,
            sequence_length=seq_len,
            batch_size=1,
            ttft_ms=round(ttft, 2),
            tps=round(tps, 2),
            total_latency_ms=round(total_time, 2),
            peak_vram_gb=round(peak_vram, 3),
            avg_vram_gb=round(current_vram, 3),
        )

    def run_suite(
        self,
        prompts: list[str],
        max_new_tokens: int = 256,
    ) -> list[BenchmarkResult]:
        """여러 프롬프트에 대해 벤치마크를 수행하고 결과를 반환한다."""
        results = []
        for prompt in prompts:
            # GPU 캐시 정리로 측정 간 간섭 최소화
            torch.cuda.empty_cache()
            result = self.measure_generation(prompt, max_new_tokens)
            results.append(result)
        return results


def compare_models(
    results_a: list[BenchmarkResult],
    results_b: list[BenchmarkResult],
) -> dict:
    """두 모델의 벤치마크 결과를 비교 요약한다."""
    def avg(lst, field):
        values = [getattr(r, field) for r in lst]
        return sum(values) / len(values) if values else 0

    return {
        "model_a": results_a[0].model_name if results_a else "unknown",
        "model_b": results_b[0].model_name if results_b else "unknown",
        "ttft_ratio": avg(results_b, "ttft_ms") / max(avg(results_a, "ttft_ms"), 0.01),
        "tps_ratio": avg(results_b, "tps") / max(avg(results_a, "tps"), 0.01),
        "vram_ratio": avg(results_b, "peak_vram_gb") / max(avg(results_a, "peak_vram_gb"), 0.01),
        "note": "ratio > 1 means model_b uses more of the metric than model_a",
    }

"""
요청 특성에 따라 Transformer/Mamba 모델을 동적으로 라우팅하는 서비스.

단순한 시퀀스 길이 기준부터 시작해서, 품질 피드백을 기반으로
라우팅 정책을 점진적으로 개선한다.
"""
from dataclasses import dataclass
from enum import Enum
from typing import Protocol


class ModelType(Enum):
    TRANSFORMER = "transformer"
    MAMBA = "mamba"


@dataclass
class RoutingDecision:
    model: ModelType
    reason: str
    confidence: float  # 0.0 ~ 1.0


@dataclass
class RequestFeatures:
    sequence_length: int
    task_type: str          # "chat", "summarization", "retrieval", "coding"
    latency_budget_ms: int
    gpu_memory_gb: float
    requires_in_context_learning: bool = False


class ModelRouter:
    """규칙 기반 모델 라우터.

    초기 도입 시에는 명시적 규칙으로 시작하고,
    A/B 테스트 결과를 반영해 점진적으로 룰을 갱신한다.
    """

    # 시퀀스 길이 임계값: 이 이상이면 Mamba가 VRAM 효율적
    LONG_CONTEXT_THRESHOLD = 32_000
    # 매우 짧은 시퀀스에서는 Transformer 오버헤드가 미미
    SHORT_CONTEXT_THRESHOLD = 2_000

    def route(self, features: RequestFeatures) -> RoutingDecision:
        # 규칙 1: in-context learning이 필요하면 Transformer
        if features.requires_in_context_learning:
            return RoutingDecision(
                model=ModelType.TRANSFORMER,
                reason="In-context learning requires attention mechanism",
                confidence=0.9,
            )

        # 규칙 2: 긴 시퀀스 + 제한된 GPU 메모리 -> Mamba
        if (features.sequence_length > self.LONG_CONTEXT_THRESHOLD
                and features.gpu_memory_gb <= 24):
            return RoutingDecision(
                model=ModelType.MAMBA,
                reason=f"Seq len {features.sequence_length} exceeds {self.LONG_CONTEXT_THRESHOLD} "
                       f"with only {features.gpu_memory_gb}GB VRAM",
                confidence=0.85,
            )

        # 규칙 3: 긴 시퀀스 + 충분한 GPU -> 태스크 유형으로 판단
        if features.sequence_length > self.LONG_CONTEXT_THRESHOLD:
            if features.task_type in ("summarization", "retrieval"):
                return RoutingDecision(
                    model=ModelType.MAMBA,
                    reason="Long-context summarization/retrieval benefits from linear scaling",
                    confidence=0.7,
                )

        # 규칙 4: 짧은 시퀀스에서는 Transformer 생태계 이점 활용
        if features.sequence_length < self.SHORT_CONTEXT_THRESHOLD:
            return RoutingDecision(
                model=ModelType.TRANSFORMER,
                reason="Short sequences: Transformer ecosystem (vLLM, TensorRT) is more optimized",
                confidence=0.8,
            )

        # 기본값: Transformer (생태계 성숙도 우위)
        return RoutingDecision(
            model=ModelType.TRANSFORMER,
            reason="Default: Transformer ecosystem maturity",
            confidence=0.5,
        )

def compare_tokenization(text: str, tokenizer_a, tokenizer_b) -> dict:
    """두 토크나이저의 분절 차이를 비교한다.

    모델 교체 시 토큰 수 차이가 비용과 품질에 미치는 영향을
    사전에 파악하기 위한 도구.
    """
    tokens_a = tokenizer_a.encode(text)
    tokens_b = tokenizer_b.encode(text)

    return {
        "text_length_chars": len(text),
        "tokenizer_a_count": len(tokens_a),
        "tokenizer_b_count": len(tokens_b),
        "token_ratio": len(tokens_b) / len(tokens_a) if tokens_a else 0,
        "cost_impact": f"Model B uses {len(tokens_b)/len(tokens_a)*100:.1f}% tokens vs Model A"
        if tokens_a else "N/A",
        # 한국어 등 비라틴 텍스트에서 차이가 큼
        "sample_tokens_a": tokenizer_a.convert_ids_to_tokens(tokens_a[:20]),
        "sample_tokens_b": tokenizer_b.convert_ids_to_tokens(tokens_b[:20]),
    }

# CI에서 모델 회귀 테스트를 자동화하는 GitHub Actions 예시
name: model-regression-gate
on:
  pull_request:
    paths:
      - 'requirements*.txt'
      - 'Dockerfile*'
      - 'configs/model_*.yaml'
  schedule:
    - cron: '0 6 * * 1' # 매주 월요일 06:00 UTC

jobs:
  regression-test:
    runs-on: [self-hosted, gpu, a100]
    steps:
      - uses: actions/checkout@v4

      - name: Install dependencies
        run: pip install -r requirements-inference.txt

      - name: Run quality benchmark
        run: |
          python scripts/benchmark.py \
            --models transformer,mamba \
            --suite production_sample \
            --output results/regression_$(date +%Y%m%d).json

      - name: Check quality gate
        run: |
          python scripts/check_regression.py \
            --results results/regression_$(date +%Y%m%d).json \
            --min-quality 0.82 \
            --max-p95-latency-ms 900 \
            --max-vram-gb 24

      - name: Upload results
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: regression-results
          path: results/

[시작]
  |
  v
시퀀스 길이가 주로 32K 이상인가? ──No──> Transformer 유지
  |                                        (생태계 이점이 압도적)
 Yes
  |
  v
GPU 메모리가 충분한가 (>=80GB)? ──Yes──> Transformer + chunked prefill
  |                                        (품질 확보가 우선이면)
  No
  |
  v
in-context learning이 핵심인가? ──Yes──> Transformer (작은 모델 + 긴 컨텍스트)
  |                                        또는 Jamba 같은 하이브리드
  No
  |
  v
품질 요구가 "Transformer 수준"인가? ──Yes──> 하이브리드 (Mamba + Attention layers)
  |
  No (약간의 품질 트레이드오프 허용)
  |
  v
Mamba 도입 후 A/B 테스트로 검증