필사 모드: 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:
        """测量单条 prompt 的生成性能。"""
        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：到生成第一个 token 为止的时间
        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]:
        """对多条 prompt 执行基准测试并返回结果。"""
        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:
    """比较两个分词器的切分差异。

    用于在更换模型之前，预先掌握 token 数量差异
    对成本和质量造成的影响。
    """
    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 测试验证