"""
=============================================================================
CS349E · 第 6 讲 · 02_selective_batching.py —— 选择性批处理：长短不一的请求怎样放进一个批
=============================================================================

【这个脚本要回答的问题】
    迭代级调度每次迭代选出来的请求处在不同阶段：有的刚到（要处理全部输入，Orca 叫 initiation），有的在生成
    （每次 1 个 token，Orca 叫 increment），而且已有的上下文长度各不相同。普通的批处理要求每个请求的张量形状一样，
    这些请求没法直接拼成 [B, L, d]。Orca 的办法——选择性批处理（selective batching）：
        线性层、归一化、残差这些"逐 token"的运算：把所有 token 拍平成 [ΣL, d]，一起算；
        注意力（需要知道哪些 token 属于同一个请求）：按请求拆开，各算各的，再拼回去。

      演示 1：Orca 图 4/图 5 的例子（4 个请求、7 个 token）在第 2 讲的 numpy Transformer 上真算一遍，
              和"每个请求单独算"逐位比较。
      演示 2：为什么不能用普通的批处理：形状对不上；硬补齐（padding）要算多少。
      演示 3：账本：拆开注意力的代价是什么？一次选择性批处理 vs 4 次单独前向，逐类比较 FLOPs、读权重、读 KV。
      演示 4：放到真实尺寸（GPT-3 13B、A100）：一次迭代的 roofline 时间；以及"注意力一个请求一个内核"的启动开销。

【运行】python3 02_selective_batching.py      （< 1 秒）
【依赖】numpy；本目录的 decode_lab；../lecture02 的 transformer_lab；../lecture01 的 console。
"""

import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)

from decode_lab import A100_40G, GPT3_13B, StepModel, ragged_forward  # noqa: E402
from console import fmt_bytes, fmt_num, fmt_time, print_header, print_observations, print_table  # noqa: E402
from transformer_lab import TINY, Counter, KVCache, forward, init_params  # noqa: E402

PARAMS = init_params(TINY, seed=0)

# Orca 图 4/图 5 的 4 个请求（token 编号是随便取的整数）：
#   x1：输入 2 个 token，已经生成了 x13，这次喂进 x14（increment，KV cache 里有 x11 x12 x13）
#   x2：输入 1 个 token x21，这次喂进 x22（increment，KV cache 里有 x21）
#   x3：刚到，输入 2 个 token（initiation）
#   x4：刚到，输入 3 个 token（initiation）
HISTORY = {"x1": [11, 12, 13], "x2": [21], "x3": [], "x4": []}
NEW = {"x1": [14], "x2": [22], "x3": [31, 32], "x4": [41, 42, 43]}


def fresh_caches():
    """为 4 个请求各建一个 KV cache，并把它们之前的 token（HISTORY）先算进去。"""
    caches = {}
    for name, hist in HISTORY.items():
        c = KVCache(TINY, 1)
        if hist:
            forward(PARAMS, TINY, np.array([hist]), cache=c)
        caches[name] = c
    return caches


# -----------------------------------------------------------------------------
def demo1_orca_example():
    print_header("演示 1：Orca 图 5 的例子——一次迭代里 4 个请求、7 个 token")
    caches = fresh_caches()
    names = list(NEW)
    chunks = [np.array(NEW[n]) for n in names]
    print("拍平后的输入：" + " ".join(f"{n}:{NEW[n]}" for n in names) + f" → 张量 [1, {sum(len(c) for c in chunks)}, d={TINY.d_model}]")
    outs = ragged_forward(PARAMS, TINY, chunks, [caches[n] for n in names])
    ref_caches = fresh_caches()
    rows = []
    worst = 0.0
    for n, o in zip(names, outs):
        ref = forward(PARAMS, TINY, np.array([NEW[n]]), cache=ref_caches[n])[0]
        diff = float(np.max(np.abs(o - ref)))
        worst = max(worst, diff)
        rows.append([n, "increment" if HISTORY[n] else "initiation", len(HISTORY[n]), len(NEW[n]),
                     f"[{len(NEW[n])}, {len(HISTORY[n]) + len(NEW[n])}]", f"{diff:.1e}"])
    print_table(["请求", "阶段", "已有 KV", "新 token", "注意力分数矩阵形状", "和单独计算的最大差"], rows)
    print_observations([
        f"选择性批处理的 logits 和\"每个请求单独前向\"的最大差 {worst:.0e}：只是浮点加法顺序的差别，计算本身完全一样。",
        "线性层看到的是一个 [7, d] 的矩阵，不知道也不需要知道这 7 行来自 4 个请求；"
        "只有注意力需要知道\"谁和谁是同一个请求\"——4 个请求的分数矩阵形状各不相同（1×4、1×2、2×2、3×3）。",
    ])


# -----------------------------------------------------------------------------
def demo2_padding():
    print_header("演示 2：为什么普通的批处理不行（Orca §3 C2 的三种情况）")
    rows = []
    for n in NEW:
        rows.append([n, f"[{len(NEW[n])}, d]", f"[{len(HISTORY[n])}, G, dh]"])
    print_table(["请求", "这次输入的形状", "KV cache 的形状"], rows)
    lens = [len(v) for v in NEW.values()]
    ctx = [len(HISTORY[n]) + len(NEW[n]) for n in NEW]
    B, Lmax, Smax = len(lens), max(lens), max(ctx)
    print(f"\n硬拼成 [B, L, d] 要补齐到 [{B}, {Lmax}, d] = {B * Lmax} 个 token 位置（真正有用的 {sum(lens)} 个）；"
          f"注意力要补齐到 [{B}, {Lmax}, {Smax}] 的分数矩阵 = {B * Lmax * Smax} 个分数（有用的 {sum(a * b for a, b in zip(lens, ctx))} 个）。")
    print("而且补齐之后还要额外的掩码，告诉注意力\"哪些位置是补的\"，KV cache 的写入位置也各不相同。")
    print_observations([
        "三种不能直接批处理的情况：(1) 都在 initiation 但输入长度不同（x3、x4）；(2) 都在 increment 但位置不同（x1、x2）；"
        "(3) 一个 initiation 一个 increment（x1、x3）。批越大，碰巧\"形状完全一样\"的概率越小。",
        "选择性批处理把\"批\"这一维去掉：线性层只关心 token 的总数 ΣL，补齐的浪费变成 0。",
    ])


# -----------------------------------------------------------------------------
def demo3_ledger():
    print_header("演示 3：账本——一次选择性批处理 vs 4 次单独前向（第 2 讲的 TINY 模型，按 BF16 每个数 2 字节计）")
    caches = fresh_caches()
    names = list(NEW)
    c_sel = Counter()
    ragged_forward(PARAMS, TINY, [np.array(NEW[n]) for n in names], [caches[n] for n in names], counter=c_sel)
    caches = fresh_caches()
    c_sep = Counter()
    for n in names:
        forward(PARAMS, TINY, np.array([NEW[n]]), cache=caches[n], counter=c_sep)
    lin = ("q_proj", "k_proj", "v_proj", "o_proj", "mlp_gate", "mlp_up", "mlp_down", "lm_head")
    rows = []
    for label, ops in (("线性层（含 lm_head）", lin), ("注意力核心", ("attn_core",)), ("写 KV", ("kv_write",))):
        rows.append([label,
                     fmt_num(c_sel.total("flops", ops)), fmt_num(c_sep.total("flops", ops)),
                     fmt_bytes(c_sel.total("w", ops)), fmt_bytes(c_sep.total("w", ops)),
                     fmt_bytes(c_sel.total("kv", ops)), fmt_bytes(c_sep.total("kv", ops))])
    print_table(["运算", "FLOPs 选择性", "FLOPs 单独", "读权重 选择性", "读权重 单独", "读写 KV 选择性", "读写 KV 单独"], rows)
    print("\n（单独前向时 lm_head 对每个新 token 都算——和选择性批处理一样，所以 FLOPs 相同。）")
    print_observations([
        "FLOPs 完全一样：选择性批处理不会多算也不会少算。",
        "读权重：选择性批处理读 1 遍，4 次单独前向读 4 遍——批处理在 decode 时省的就是这个（第 1 讲 §4.2）。",
        "注意力核心读的是各个请求自己的 KV，本来就不能共享：拆开算，读的字节一点没多。"
        "这正是 Orca 的论点：注意力没有权重，不对它做批处理，几乎不损失效率。",
    ])


# -----------------------------------------------------------------------------
def demo4_real_size():
    print_header("演示 4：真实尺寸——GPT-3 13B、A100 40GB 上的一次迭代")
    step = StepModel(GPT3_13B, A100_40G)
    m = GPT3_13B
    rng = np.random.default_rng(0)
    ctx = rng.integers(100, 600, 32)
    rows = []
    t_dec_only = None
    for label, prompts in (("32 个请求在生成", []), ("32 个请求在生成 + 2 个新请求（输入 300、450）", [300, 450])):
        attn = m.attn_flops_per_token(1) * sum(L * (L + 1) / 2 for L in prompts)
        t_sel = step.mixed_time(32, int(ctx.sum()), sum(prompts), attn)
        t_sep = sum(step.decode_time(1, int(c)) for c in ctx) + sum(step.prefill_time([L]) for L in prompts)
        n_req = 32 + len(prompts)
        launches = m.n_layers * n_req
        cells = [label, fmt_time(t_sel), f"{fmt_time(t_sep)}（{t_sep / t_sel:.0f}×）"]
        for launch in (3e-6, 10e-6):
            cells.append(f"+{launches * launch / t_sel:.0%}")
        rows.append(cells)
        if not prompts:
            t_dec_only = (t_sel, launches)
    print_table(["这次迭代里有", "选择性批处理", "每个请求单独前向", "注意力逐请求启动 3 µs", "10 µs"], rows)
    print(f"\n（32 个生成中的请求上下文长度 U(100, 600)，总长 {int(ctx.sum()):,}；每步固定开销 1 ms；"
          "后两列 = 每层给每个请求单独启动一次注意力内核，额外的启动时间占这次迭代的比例，启动开销是假设值。）")
    t0, n_l = t_dec_only
    print_observations([
        "每个请求单独前向要几倍到几十倍的时间：权重被反复读。纯 decode 的迭代差得最多（每个请求单独读一遍 25.7 GB）；"
        "有新请求时，它们的输入本身就够把 GPU 算满，差距变小。",
        f"注意力按请求拆开以后，如果真的每层给每个请求启动一个内核：纯 decode 的迭代（{fmt_time(t0)}）里 "
        f"{n_l} 次启动就多出 {n_l * 3e-6 / t0:.0%}–{n_l * 10e-6 / t0:.0%} 的时间。所以 Orca 把各个请求的注意力内核"
        "\"拼\"成一次启动（不同请求的线程块放进同一个内核，§5）。",
        "今天的引擎（vLLM、SGLang、TensorRT-LLM）用专门的\"变长\"注意力内核（FlashAttention varlen、FlashInfer）："
        "一次启动处理一批长短不一的请求，KV 按块存放（第 7 讲）——这就是选择性批处理的现代形态。",
    ])


if __name__ == "__main__":
    demo1_orca_example()
    demo2_padding()
    demo3_ledger()
    demo4_real_size()