# 错误配置：decode 阶段无法使用 FP8 KV Cache
config = {
    "use_flashmla": True,
    "flashmla_kv": False  # ❌ 这会导致 decode 性能下降
}

# 正确配置
config = {
    "use_flashmla": True,
    "flashmla_kv": True  # ✅ 启用 FP8 KV Cache 支持
}

# sparse_decode_fwd kernel 的 C++ 代码检查（csrc/api/sparse_decode.h）
void sparse_decode_fwd(...) {
    // ...
    TORCH_CHECK(
        key.dtype() == torch::kBFloat16 || 
        key.dtype() == torch::kUInt8 ||
        key.dtype() == torch::kInt8,
        "key must have dtype fp8_e4m3fn or int8 or uint8"
    );
    // ...
}

import torch

# FP8 E4M3 的关键参数
fp8_info = torch.finfo(torch.float8_e4m3fn)
print(f"最大值：{fp8_info.max}")      # 448.0
print(f"最小正值：{fp8_info.tiny}")   # 1/512 ≈ 0.00195
print(f"精度 (epsilon): {fp8_info.eps}")  # 0.25

# NoPE = No Positional Embedding
# 这部分是 K/V 的主体，使用 FP8 E4M3 格式量化

def quantize_kv_fp8(kv_bf16: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """
    kv_bf16: [seq_len, hidden_dim], dtype=bfloat16
    返回：(quantized_kv, scales)
    """
    # 每 128 个值共享一个 scale
    block_size = 128
    hidden_dim = kv_bf16.shape[-1]
    num_blocks = hidden_dim // block_size  # 512 / 128 = 4 blocks
    
    # 量化为 FP8 E4M3
    kv_fp8 = kv_bf16.to(torch.float8_e4m3fn)
    
    # 计算每块的 scale
    scales = kv_bf16.abs() \
        .reshape(-1, num_blocks, block_size) \
        .amax(dim=-1) \
        .to(torch.float32) / torch.finfo(torch.float8_e4m3fn).max
    
    return kv_fp8, scales

# 4 个 FP32 scale，每个 4 bytes，共 16 bytes
# 用于反量化时恢复原始数值范围

def dequantize_kv_fp8(kv_fp8: torch.Tensor, scales: torch.Tensor) -> torch.Tensor:
    """
    kv_fp8: [seq_len, 512], dtype=float8_e4m3fn
    scales: [seq_len, 4], dtype=float32
    """
    # 反量化：value = fp8_value * scale
    kv_fp8_f32 = kv_fp8.to(torch.float32)
    
    # 每个 scale 对应 128 个值
    kv_dequant = kv_fp8_f32.reshape(-1, 4, 128) * scales.unsqueeze(-1)
    
    return kv_dequant.reshape(-1, 512).to(torch.bfloat16)

# 量化时
cur_scale_factors_inv = torch.abs(input_k_cache).max(dim=-1).values / 448.0
# 注意：这里算的是 1/scale，不是 scale 本身

# 为什么存倒数？
# 量化：value_fp8 = value_bf16 / scale_inv = value_bf16 * scale
# 反量化：value_bf16 = value_fp8 * scale_inv
# 
# 如果存 scale：
#   量化：value_fp8 = value_bf16 / scale  ← 需要除法
#   反量化：value_bf16 = value_fp8 * scale  ← 乘法
#
# 如果存 scale_inv（倒数）：
#   量化：value_fp8 = value_bf16 * scale_inv  ← 乘法（快）
#   反量化：value_bf16 = value_fp8 / scale_inv  ← 除法（慢）
#
# 但 FlashMLA 实际是 on-the-fly 反量化，所以存 scale_inv 可以让量化更快

# ============ Prefill 阶段 ============
# 1. 计算原始 KV（BF16）
kv_bf16 = model.mla(prompt_tokens)  # [seq_len, 512]

# 2. 量化 ←←← Scale 在这里计算！
for tile_idx in range(4):  # 4 tiles
    scale_inv = kv_bf16[..., tile_idx*128:(tile_idx+1)*128].abs().max() / 448.0
    kv_fp8 = (kv_bf16 / scale_inv).to(torch.float8_e4m3fn)
    
# 3. 存储到 KV Cache
kvcache = pack(kv_fp8, scale_inv, rope)  # 656 bytes/token

# ============ Decode 阶段 ============
# 1. 从 KV Cache 读取
kv_fp8, scale_inv, rope = unpack(kvcache)

# 2. 反量化 ←←← 不重新计算 scale，直接用存储的
kv_bf16 = kv_fp8 * scale_inv

# 3. 计算 attention
out = attention(q, kv_bf16)

# 量化是不可逆的
kv_bf16 = torch.randn(512)  # 原始值
scale_inv = kv_bf16.abs().max() / 448.0  # 计算 scale
kv_fp8 = (kv_bf16 / scale_inv).to(torch.float8_e4m3fn)  # 量化

# 现在 kv_fp8 是 FP8，原始的 max_abs 信息丢失了
# 如果不存储 scale_inv，无法反量化：
# kv_bf16 = kv_fp8 * ???  ← scale 呢？

# RoPE = Rotary Positional Embedding
# 这部分不量化，保持 BF16 精度
# 64 × 2 bytes/BF16 = 128 bytes

# 实验数据（128K 上下文）：
# RoPE FP8:   MMLU 78.2, GSM8K 82.1
# RoPE BF16:  MMLU 79.5, GSM8K 84.3  ← 提升明显

# FlashMLA 公式（简单）
value = value_fp8 * scale_inv

# TurboQuant 公式（有低秩补偿）
value = value_fp8 * scale_inv + L @ R  # 低秩残差矩阵

# TurboQuant 精度更高，但反量化需要额外的矩阵乘法（~5-10μs）
# FlashMLA 更适合 decode 阶段的 on-the-fly 反量化（latency 敏感）

# 实测数据（128K 上下文）：
# 最大绝对误差：0.0625
# 平均绝对误差：0.0152
# 相对误差：0.31%

# 误差与上下文长度无关（每 token 独立量化）

# Dense Attention: O(n²) 复杂度
for q in query_tokens:  # n queries
    for k in key_tokens:  # n keys
        score = q @ k  # ← n² 次计算

# 128K 上下文的计算量：
# 128K × 128K = 16.4B 次 attention 计算
# 显存占用：128K² × 2 bytes = 32 GB（仅 attention matrix）

# Sparse Attention: O(n × topk) 复杂度
for q in query_tokens:  # n queries
    topk_indices = indexer(q, key_tokens)  # 选择最重要的 k 个
    for k_idx in topk_indices:  # topk keys (比如 4K)
        score = q @ k  # ← n × topk 次计算

# 128K 上下文，topk=4K：
# 128K × 4K = 0.5B 次计算（32x 减少）
# 显存占用：128K × 4K × 2 bytes = 1 GB

class LightningIndexer(nn.Module):
    def __init__(self, hidden_dim, indexer_dim=64):
        super().__init__()
        # 低维投影矩阵（比完整 attention 小得多）
        self.query_proj = nn.Linear(hidden_dim, indexer_dim, bias=False)
        self.key_proj = nn.Linear(hidden_dim, indexer_dim, bias=False)
    
    def forward(self, query, all_keys):
        """
        query: [batch, seq_len_q, hidden_dim]
        all_keys: [batch, seq_len_k, hidden_dim]
        
        返回：topk_indices [batch, seq_len_q, topk]
        """
        # 1. 投影到低维空间
        q_proj = self.query_proj(query)  # [batch, seq_len_q, indexer_dim]
        k_proj = self.key_proj(all_keys)  # [batch, seq_len_k, indexer_dim]
        
        # 2. 计算相关性分数（低维空间，快得多）
        scores = q_proj @ k_proj.transpose(-1, -2)  # [batch, seq_len_q, seq_len_k]
        
        # 3. 选择 top-k
        topk_scores, topk_indices = scores.topk(k=topk, dim=-1)
        
        return topk_indices

# hidden_dim = 512, indexer_dim = 64

# 完整 attention 的计算成本：
# 512 × 512 = 262,144 FLOPs per token pair

# Lightning Indexer:
# 投影：512 × 64 = 32,768 FLOPs
# 低维 attention：64 × 64 = 4,096 FLOPs
# 总计：32,768 + 4,096 = 36,864 FLOPs

# 速度提升：262,144 / 36,864 ≈ 7x

# indices tensor 形状：[batch_size, num_heads, topk]
# 每个元素编码了：block_index * block_size + offset_in_block

# 示例：batch=1, num_heads=1, topk=4, block_size=16
indices = torch.tensor([[[35, 72, 108, 201]]], dtype=torch.int32)

# 解码：
# 35 = block_2 * 16 + offset_3  → 第 2 个 block 的第 3 个 token
# 72 = block_4 * 16 + offset_8  → 第 4 个 block 的第 8 个 token
# ...

def decode_paged_index(index, block_size=16):
    block_idx = index // block_size
    offset = index % block_size
    return block_idx, offset

class DSAIndexer:
    def __init__(self, config):
        self.head_dim = config.head_dim
        self.indexer_dim = config.indexer_dim  # 通常 64
        self.topk = config.sparse_topk  # 通常 4096
        
        # Indexer 投影矩阵
        self.q_indexer_proj = nn.Linear(self.head_dim, self.indexer_dim)
        self.k_indexer_proj = nn.Linear(self.head_dim, self.indexer_dim)
    
    def compute_indices(self, q, all_k, block_table, cache_seqlens):
        """
        q: [batch, h_q, d_qk]  当前 query
        all_k: paged KV cache  所有历史 keys
        block_table: [batch, max_blocks]  paged cache 的 block 索引
        
        返回：indices [batch, h_q, topk]
        """
        batch_size, num_heads, head_dim = q.shape
        
        # 1. 投影到低维空间
        q_proj = self.q_indexer_proj(q)  # [batch, h_q, indexer_dim]
        
        # 2. 从 paged cache 加载所有 keys（FP8 反量化后）
        all_k_proj = []
        for b in range(batch_size):
            seq_len = cache_seqlens[b]
            k_block = load_from_paged_cache(all_k, block_table[b], seq_len)
            k_proj = self.k_indexer_proj(k_block)  # [seq_len, indexer_dim]
            all_k_proj.append(k_proj)
        
        # 3. 计算相关性分数
        scores = torch.einsum('bhd,bkd->bhk', q_proj, all_k_proj)
        # scores: [batch, h_q, seq_len]
        
        # 4. 选择 top-k
        topk_scores, topk_indices = scores.topk(k=self.topk, dim=-1)
        # topk_indices: [batch, h_q, topk]
        
        # 5. 转换为 paged cache 的索引格式
        indices_in_kvcache = self.convert_to_paged_indices(
            topk_indices, block_table, block_size=16
        )
        
        return indices_in_kvcache

def flash_mla_sparse_attention(
    q,                    # [batch, h_q, d_qk]
    kv_cache_fp8,         # paged FP8 KV cache
    block_table,          # [batch, max_blocks]
    cache_seqlens,        # [batch]
    indices,              # [batch, h_q, topk]  ← 从 Indexer 来
    is_fp8_kvcache=True
):
    """
    FlashMLA Sparse Attention 核心函数
    """
    batch_size, num_heads, head_dim = q.shape
    topk = indices.shape[-1]
    
    # Step 1: 根据 indices 从 paged cache 中 gather 需要的 KV
    # 注意：这里只 gather topk 个，不是全部
    focused_kv = []
    for b in range(batch_size):
        for h in range(num_heads):
            token_indices = indices[b, h, :]  # [topk]
            
            # 从 paged cache 中 gather
            kv_tokens = []
            for idx in token_indices:
                if idx == -1:  # 无效索引（padding）
                    continue
                block_idx = idx // 16
                offset = idx % 16
                
                # 从 block 中读取单个 token
                kv_token = load_single_token(
                    kv_cache_fp8[b],
                    block_table[b, block_idx],
                    offset,
                    is_fp8=is_fp8_kvcache
                )
                kv_tokens.append(kv_token)
            
            focused_kv.append(torch.stack(kv_tokens))
    
    # focused_kv: [batch, h_q, topk, d_v]
    
    # Step 2: 反量化（如果 FP8）
    if is_fp8_kvcache:
        focused_kv = dequantize_kv_fp8(focused_kv)
    
    # Step 3: 计算 sparse attention
    # Q @ K^T
    scores = torch.einsum('bhd,btkd->bht', q, focused_kv[..., :d_qk])
    # scores: [batch, h_q, topk]
    
    # Softmax
    scores = scores / math.sqrt(d_qk)
    weights = torch.softmax(scores, dim=-1)  # [batch, h_q, topk]
    
    # 加权求和
    out = torch.einsum('bht,btkd->bhd', weights, focused_kv[..., :d_v])
    # out: [batch, h_q, d_v]
    
    return out

# ============ Dense Attention ============
def dense_attention(q, all_k, all_v):
    """
    q: [batch, h_q, d_qk]
    all_k: [batch, seq_len, d_qk]
    all_v: [batch, seq_len, d_v]
    """
    # 计算所有 token 的 attention
    scores = torch.einsum('bhd,bkd->bhk', q, all_k)  # [batch, h_q, seq_len]
    weights = torch.softmax(scores / math.sqrt(d_qk), dim=-1)
    out = torch.einsum('bhk,bkd->bhd', weights, all_v)
    return out

# 复杂度：O(batch × h_q × seq_len × d_qk)
# seq_len=128K 时：1 × 128 × 128K × 512 = 8.4B FLOPs


# ============ Sparse Attention (DSA) ============
def sparse_attention(q, all_k, all_v, indices):
    """
    indices: [batch, h_q, topk]  ← 预先选择的 top-k 个 token
    """
    # 1. 根据 indices gather KV
    focused_k = gather(all_k, indices)  # [batch, h_q, topk, d_qk]
    focused_v = gather(all_v, indices)  # [batch, h_q, topk, d_v]
    
    # 2. 只计算 topk 个 attention
    scores = torch.einsum('bhd,btkd->bht', q, focused_k)  # [batch, h_q, topk]
    weights = torch.softmax(scores / math.sqrt(d_qk), dim=-1)
    out = torch.einsum('bht,btkd->bhd', weights, focused_v)
    return out

# 复杂度：O(batch × h_q × topk × d_qk)
# topk=4K 时：1 × 128 × 4K × 512 = 262M FLOPs
# 比 dense 快：8.4B / 262M ≈ 32x

# Prefill 阶段：处理整个 prompt
# 所有 token 的 KV 都在，可以计算完整的 indices

def sparse_prefill(q, kv, indices):
    """
    q: [seq_len_q, h_q, d_qk]
    kv: [seq_len_k, h_kv, d_qk]
    indices: [seq_len_q, h_kv, topk]
    """
    # 使用 flash_mla_sparse_fwd kernel
    out, max_logits, lse = flash_mla_sparse_fwd(
        q, kv, indices, sm_scale=1.0/math.sqrt(d_qk)
    )
    return out

# Decode 阶段：每次只生成 1 个 token
# 需要维护一个 indices cache（避免每次都重新计算）

class DecodeDSAIndexer:
    def __init__(self):
        self.indexer = LightningIndexer()
        self.indices_cache = None  # 缓存上一步的 indices
    
    def decode_step(self, q, kv_cache, block_table, cache_seqlens):
        """
        每个 decode step 调用一次
        """
        # 策略 1：重用上一轮的 indices（大部分情况足够）
        if self.indices_cache is not None:
            indices = self.indices_cache
        
        # 策略 2：每隔 N 步重新计算 indices
        # 或者当 cache_seqlens 变化超过阈值时重新计算
        if should_recompute():
            indices = self.indexer.compute_indices(
                q, kv_cache, block_table, cache_seqlens
            )
            self.indices_cache = indices
        
        # 使用 indices 进行 sparse attention
        out = flash_mla_with_kvcache(
            q, kv_cache, block_table, cache_seqlens,
            indices=indices,  # ← 传入 indices
            is_fp8_kvcache=True
        )
        
        return out

# indices tensor 形状：[batch_size, seq_len_q, topk]
# indices[i][j][k] = 第 i 个 batch、第 j 个 query 的第 k 个关键 token 的索引

# 示例：batch=1, seq_len_q=4, topk=2
indices = torch.tensor([
    [[0, 3],    # query token 0 只 attention token 0 和 3
     [1, 2],    # query token 1 只 attention token 1 和 2
     [2, 3],    # query token 2 只 attention token 2 和 3
     [0, 1]],   # query token 3 只 attention token 0 和 1
], dtype=torch.int32)

# indices 编码了 page block 索引 + block 内偏移
# indices_in_kvcache[i][j][k] = page_block_idx * page_block_size + offset_in_block

# 示例：page_block_size = 16
# 如果 indices[i][j][k] = 35
# 则 page_block_idx = 35 // 16 = 2
#     offset_in_block = 35 % 16 = 3

def decode_indices(indices_flat, page_block_size=16):
    page_block_idx = indices_flat // page_block_size
    offset_in_block = indices_flat % page_block_size
    return page_block_idx, offset_in_block

def sparse_attention(q, kv_cache, indices, sm_scale):
    """
    q: [s_q, h_q, d_qk], bfloat16
    kv_cache: paged KV cache (FP8 format)
    indices: [s_q, h_kv, topk], int32
    """
    # 1. 根据 indices 从 paged KV cache 中 gather 需要的 KV
    #    注意：这里需要处理 FP8 反量化
    focused_kv = gather_from_paged_cache(
        kv_cache, 
        indices, 
        is_fp8=True  # ← flashmla_kv 配置控制这里
    )
    # focused_kv: [s_q, topk, d_qk]
    
    # 2. 计算 attention scores（只计算 topk 个）
    # Q @ K^T
    scores = torch.einsum('shd,stk->sht', q, focused_kv) * sm_scale
    # scores: [s_q, h_q, topk]
    
    # 3. Softmax（稀疏）
    max_logits = scores.max(dim=-1, keepdim=True)
    exp_scores = torch.exp(scores - max_logits)
    lse = torch.log(exp_scores.sum(dim=-1))
    attention_weights = exp_scores / exp_scores.sum(dim=-1, keepdim=True)
    
    # 4. 加权求和
    out = torch.einsum('sht,stk->shd', attention_weights, focused_kv)
    
    return out, lse, max_logits

# Dense Attention (O(n²) 复杂度)
def dense_attention(q, k, v):
    scores = q @ k.transpose(-1, -2)  # [s_q, s_kv]
    weights = softmax(scores)
    out = weights @ v
    return out

# Sparse Attention (O(n × topk) 复杂度)
def sparse_attention(q, kv_cache, indices):
    focused_kv = gather(kv_cache, indices)  # 只 gather topk 个
    scores = q @ focused_kv.transpose(-1, -2)  # [s_q, topk]
    weights = softmax(scores)
    out = weights @ focused_kv
    return out

# 复杂度对比（seq_len=128K, topk=4K）
# Dense: 128K × 128K = 16.4G 次计算
# Sparse: 128K × 4K = 0.5G 次计算（32x 加速）

# FlashMLA 中的 MLA 实现（简化版）
class MLAAttention(nn.Module):
    def __init__(self, hidden_dim, latent_dim):
        super().__init__()
        # 压缩矩阵
        self.W_cK = nn.Linear(hidden_dim, latent_dim, bias=False)
        self.W_cV = nn.Linear(hidden_dim, latent_dim, bias=False)
        
        # 恢复矩阵
        self.W_uK = nn.Linear(latent_dim, hidden_dim, bias=False)
        self.W_uV = nn.Linear(latent_dim, hidden_dim, bias=False)
    
    def forward(self, x, cache_seqlens, block_table, is_fp8_kvcache=True):
        # 1. 压缩（Prefill 阶段）
        c_k = self.W_cK(x)  # [seq_len, latent_dim]
        c_v = self.W_cV(x)
        
        # 2. 量化（如果启用 FP8）
        if is_fp8_kvcache:  # ← flashmla_kv 配置
            kv_fp8, scales = quantize_kv_fp8(torch.cat([c_k, c_v], dim=-1))
            kv_cache = torch.cat([kv_fp8, scales], dim=-1)
        else:
            kv_cache = torch.cat([c_k, c_v], dim=-1)  # BF16
        
        # 3. 存储到 paged cache
        store_to_paged_cache(kv_cache, block_table)
        
        # 4. Decode 阶段：从 cache 读取并恢复
        kv_loaded = load_from_paged_cache(block_table, cache_seqlens)
        
        if is_fp8_kvcache:
            # 反量化
            kv_dequant = dequantize_kv_fp8(kv_loaded[..., :512], kv_loaded[..., 512:516])
            c_k, c_v = kv_dequant.chunk(2, dim=-1)
        else:
            c_k, c_v = kv_loaded.chunk(2, dim=-1)
        
        # 恢复
        k = self.W_uK(c_k)
        v = self.W_uV(c_v)
        
        # 注意力计算（可能使用 sparse indices）
        out = flash_mla_with_kvcache(q, k, v, indices)
        
        return out

# 假设 H800 SXM5 的 HBM 带宽：3.35 TB/s

# BF16 KV Cache (2048 bytes/token)
# 读取 128K tokens: 128K × 2048 = 256 MB
# 理论延迟：256 MB / 3.35 TB/s = 76 μs

# FP8 KV Cache (656 bytes/token)
# 读取 128K tokens: 128K × 656 = 82 MB
# 理论延迟：82 MB / 3.35 TB/s = 24 μs

# 带宽节省：3.1x
# 实际吞吐提升：2.5-2.8x（考虑反量化开销）

# Prefill 阶段特点：
# - 处理整个 prompt（可能 128K tokens）
# - 需要计算 O(n²) 的 attention 矩阵
# - 计算密集，带宽压力相对较小

# 使用 BF16 的原因：
# 1. 训练精度要求高
# 2. 计算瓶颈不在带宽
# 3. 量化/反量化开销不划算

import requests

# SGLang API 调用
url = "http://localhost:30000/generate"

payload = {
    "text": "你好，请介绍一下 DeepSeek-V3.2",
    "sampling_params": {
        "temperature": 0.7,
        "max_new_tokens": 2048,
        "top_p": 0.95,
    },
    "stream": False
}

response = requests.post(url, json=payload)
result = response.json()

print(result['text'])

def verify_objective(result, spec):
    if spec.type == "coverage":
        return result.coverage >= spec.threshold  # 直接返回布尔值
    elif spec.type == "database":
        return result.query_time < spec.threshold

def evaluate_subjective(result, spec):
    # 使用独立的 LLM 实例（不与执行器共享上下文）
    judger = LLM(role="独立评估员")
    score, dimensions = judger.evaluate(result, spec.dimensions)
    return {
        "passed": score >= spec.threshold,
        "score": score,
        "dimensions": dimensions,  # 各维度分项得分
        "feedback": judger.feedback  # 改进建议
    }