1
# 传统 Batching 等待所有请求
2
# 短请求等待长请求完成
3

4
class StaticBatching:
5
    def process_batch(self, requests):
6
        # 等待所有请求准备完成
7
        while len(requests) < self.batch_size:
8
            time.sleep(0.1)
9
        return self._run_batch(requests)

1
class ContinuousBatching:
2
    """收到请求立即入队"""
3
    def add_request(self, request):
4
        self.pending.append(request)
5

6
    def step(self):
7
        # 动态调度
8
        running = [r for r in self.running if r.is_running()]
9
        self.running = running
10
        for req in self.pending[:self.max_batch]:
11
            req.start()
12
            self.running.append(req)
13

14
        # 处理完的请求立即返回，接收新请求
15
        return [r.result() for r in running if r.is_done()]

1
# vLLM 配置
2
engine_config:
3
  max_model_len: 4096
4
  tensor_parallel_size: 2
5
  gpu_memory_utilization: 0.9
6
  pipeline_parallel_size: 1

1
class KVCache:
2
    """KV Cache 存储"""
3
    def __init__(self):
4
        self.k_cache = torch.zeros(MAX_SEQ_LEN, HEAD_DIM)
5
        self.v_cache = torch.zeros(MAX_SEQ_LEN, HEAD_DIM)
6

7
    def update(self, position, k, v):
8
        self.k_cache[position] = k
9
        self.v_cache[position] = v
10

11
    def get(self, start, end):
12
        return self.k_cache[start:end], self.v_cache[start:end]

1
# vLLM 的分页管理
2
class PagedKVCache:
3
    """物理块管理虚拟 Cache"""
4
    def __init__(self, block_size=16):
5
        self.free_blocks = list(range(num_blocks))
6
        self.block_table = {}  # 请求 ID -> 物理块
7

8
    def allocate(self, request_id, seq_len):
9
        num_blocks = ceil(seq_len / block_size)
10
        blocks = [self.free_blocks.pop() for _ in range(num_blocks)]
11
        self.block_table[request_id] = blocks

1
# INT8 量化
2
class Int8Quantizer:
3
    def quantize(self, weights):
4
        # 激活
5
        scale = weights.abs().max() / 127
6
        quantized = (weights / scale).round().to(torch.int8)
7
        return quantized, scale
8

9
    def dequantize(self, quantized, scale):
10
        return quantized.float() * scale

1
# GPTQ 量化命令
2
python gptq.py \
3
    --model meta-llama/Llama-2-7b \
4
    --bits 4 \
5
    --dataset wikitext2 \
6
    --desc_act  # 描述性激活

1
# Activation-Aware Quantization
2
class AWQ:
3
    def find_scale(self, weights, activations):
4
        # 找最优缩放因子
5
        scales = {}
6
        for name, w in weights.items():
7
            act = get_activation(name)
8
            scale = (act.abs().mean() / w.abs().mean())
9
            scales[name] = scale
10
        return scales

1
# Tensor Parallel 配置
2
tensor_parallel:
3
  enabled: true
4
  size: 8 # 8 卡并行
5
  method: ring
6
  # 或 method: all-reduce

1
# AllReduce 通信优化
2
class AllReduceOptimizer:
3
    def all_reduce(self, tensor, op="sum"):
4
        # 异步通信
5
        # 计算与通信重叠
6
        pass

1
# 启动 vLLM 服务
2
python -m vllm.entrypoints.openai.api_server \
3
    --model meta-llama/Llama-2-7b \
4
    --tensor-parallel-size 2 \
5
    --gpu-memory-utilization 0.9 \
6
    --max-num-batched-tokens 4096

1
# 编译模型
2
trtllm-build --model meta-llama/Llama-2-7b \
3
    --output /tmp/model.engine \
4
    --tp 2

1
# 请求格式
2
import openai
3
client = openai.OpenAI(base_url="http://localhost:8000/v1")
4

5
response = client.chat.completions.create(
6
    model="llama-2-7b",
7
    messages=[{"role": "user", "content": "Hello"}],
8
    max_tokens=100
9
)