FP8 推理基础设施（五）：从 DeepGEMM 到 vLLM

import torch
import deep_gemm

# 创建 FP8 矩阵（模拟 LLM 推理中的 Q·K^T 操作）
# batch_size=2, seq_len=1024, hidden_dim=4096
batch, seq, dim = 2, 1024, 4096

# 输入矩阵 A (query): [batch, seq, dim]
A_fp32 = torch.randn(batch, seq, dim, device="cuda", dtype=torch.float32)
# 输入矩阵 B (key): [batch, dim, seq]
B_fp32 = torch.randn(batch, dim, seq, device="cuda", dtype=torch.float32)

# 转换为 FP8 (E4M3 格式)
A_fp8 = A_fp32.to(torch.float8_e4m3fn)
B_fp8 = B_fp32.to(torch.float8_e4m3fn)

# DeepGEMM FP8 GEMM 调用
# m=batch*seq (总 token 数), n=seq (输出序列长度), k=dim (hidden 维度)
m, n, k = batch * seq, seq, dim
C_fp8 = deep_gemm.fp8_gemm(
    A=A_fp8,
    B=B_fp8,
    m=m, n=n, k=k,
    # per-token 缩放因子（对 A）
    scale_a=torch.ones(m, device="cuda", dtype=torch.float32),
    # per-channel 缩放因子（对 B）
    scale_b=torch.ones(n, device="cuda", dtype=torch.float32),
)

# 输出转回 FP32 进行后续计算
C_fp32 = C_fp8.to(torch.float32)
print(f"输出形状: {C_fp32.shape}")  # [2048, 1024]

import torch
import time
import deep_gemm

def benchmark_gemm(m, n, k, iterations=100):
    """对比 DeepGEMM FP8 与 cuBLAS FP16 的性能"""
    
    # 准备 FP8 数据
    A_fp8 = torch.randn(m, k, device="cuda", dtype=torch.float8_e4m3fn)
    B_fp8 = torch.randn(k, n, device="cuda", dtype=torch.float8_e4m3fn)
    
    # 准备 FP16 数据（cuBLAS 基准）
    A_fp16 = torch.randn(m, k, device="cuda", dtype=torch.float16)
    B_fp16 = torch.randn(k, n, device="cuda", dtype=torch.float16)
    
    # DeepGEMM FP8 GEMM
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(iterations):
        _ = deep_gemm.fp8_gemm(A_fp8, B_fp8, m, n, k,
                               torch.ones(m, device="cuda"),
                               torch.ones(n, device="cuda"))
    torch.cuda.synchronize()
    fp8_time = (time.perf_counter() - start) / iterations
    
    # cuBLAS FP16 GEMM
    torch.cuda.synchronize()
    start = time.perf_counter()
    for _ in range(iterations):
        _ = torch.mm(A_fp16, B_fp16)
    torch.cuda.synchronize()
    fp16_time = (time.per*counter() - start) / iterations
    
    print(f"矩阵形状: {m}×{k} × {k}×{n}")
    print(f"DeepGEMM FP8: {fp8_time*1000:.2f}ms")
    print(f"cuBLAS FP16:  {fp16_time*1000:.2f}ms")
    print(f"加速比: {fp16_time/fp8_time:.2f}x")
    print(f"显存节省: 50%")

# 模拟 GPT-5.5 推理场景
# Attention: Q·K^T, batch=4, seq=2048, dim=8192
benchmark_gemm(8192, 2048, 8192)

# FFN: 第一个线性层, batch*seq=8192, dim=8192, ff_dim=22016
benchmark_gemm(8192, 22016, 8192)

import torch
import torch.nn as nn
from typing import Dict, Optional

class FP8HybridInference:
    """FP8 混合精度推理引擎
    
    策略：
    - Attention 层: FP8 (E4M3)
    - FFN 上投影: FP8
    - FFN 下投影: FP16
    - LM Head: FP16
    - Embedding: FP16
    """
    
    def __init__(self, model: nn.Module, calibration_data: torch.Tensor):
        self.model = model
        self.fp8_layers = self._identify_fp8_layers(model)
        self.scales = self._calibrate_scales(calibration_data)
    
    def _identify_fp8_layers(self, model: nn.Module) -> Dict[str, bool]:
        """识别哪些层可以用 FP8"""
        fp8_config = {}
        for name, module in model.named_modules():
            if isinstance(module, nn.Linear):
                # LM Head 和 Embedding 用 FP16
                if "lm_head" in name or "embed" in name:
                    fp8_config[name] = False
                # FFN 下投影用 FP16
                elif "down_proj" in name:
                    fp8_config[name] = False
                # 其他用 FP8
                else:
                    fp8_config[name] = True
        return fp8_config
    
    def _calibrate_scales(self, data: torch.Tensor) -> Dict[str, float]:
        """基于校准数据计算每层的 FP8 缩放因子"""
        scales = {}
        for name, module in self.model.named_modules():
            if name in self.fp8_layers and self.fp8_layers[name]:
                # 用校准数据的前向传播统计激活值范围
                with torch.no_grad():
                    # 简化版：用输入数据的 max 绝对值作为缩放因子
                    abs_max = data.abs().max().item()
                    scales[name] = abs_max / 448.0  # E4M3 最大值为 448
        return scales
    
    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
        """混合精度前向传播"""
        x = self.model.embed_tokens(input_ids)  # FP16
        
        for layer in self.model.layers:
            # Attention (FP8)
            attn_out = self._fp8_attention(layer.self_attn, x)
            x = x + attn_out
            
            # FFN 上投影 (FP8) + 下投影 (FP16)
            ff_out = self._hybrid_ffn(layer.mlp, x)
            x = x + ff_out
        
        # LM Head (FP16)
        logits = self.model.lm_head(x)
        return logits
    
    def _fp8_attention(self, attn_module, x: torch.Tensor) -> torch.Tensor:
        """FP8 注意力计算"""
        q = attn_module.q_proj(x).to(torch.float8_e4m3fn)
        k = attn_module.k_proj(x).to(torch.float8_e4m3fn)
        v = attn_module.v_proj(x).to(torch.float8_e4m3fn)
        
        # FP8 GEMM (用 DeepGEMM 或 cuBLASLt)
        # 简化版：torch.mm 自动降级到 FP16
        attn_scores = (q.to(torch.float16) @ k.transpose(-1, -2).to(torch.float16)) / (q.shape[-1] ** 0.5)
        attn_weights = torch.softmax(attn_scores, dim=-1)
        return attn_weights @ v.to(torch.float16)
    
    def _hybrid_ffn(self, mlp_module, x: torch.Tensor) -> torch.Tensor:
        """混合精度 FFN：上投影 FP8 + 下投影 FP16"""
        # 上投影：FP8
        hidden = mlp_module.up_proj(x).to(torch.float8_e4m3fn)
        hidden = mlp_module.act_fn(hidden.to(torch.float16))
        # 下投影：FP16
        return mlp_module.down_proj(hidden.to(torch.float16))


# 使用示例
# 1. 校准
calibration_data = torch.randn(512, 4096, device="cuda")
engine = FP8HybridInference(model, calibration_data)

# 2. 推理
output = engine.forward(input_ids)
print(f"输出形状: {output.shape}")  # [batch, seq, vocab]