"""
long_lab.py —— 第 19 讲的共享模块：大海捞针的玩具世界、可换位置编码和注意力模式的小 Transformer、环形注意力的模拟
=====================================================================

【这个文件要回答的问题】
    1. 一个在 128 个 token 上训练的模型，能不能在 512、1024 个 token 的"草堆"里找到一根"针"？
       位置编码（RoPE 的外推、位置插值 PI、调大 base）和继续训练怎样影响它？
    2. 近似注意力（滑动窗口、稀疏、空洞）省了计算，丢掉了什么？
    3. 环形注意力（Ring Attention）怎样把 FlashAttention 的分块推广到很多块 GPU 上？

【大海捞针的玩具世界（NeedleTask）】
    词表：<pad> <s> Q N，16 个"键" k0..k15，16 个"值" v0..v15，32 个"填充词" f0..f31。
    草堆：填充词按一条随机的马尔可夫链生成（每个填充词后面只会接 4 种填充词）——像一段有规律、但和问题无关的文字。
    针："N k v"——在草堆里某个位置插入一个键和它的值（每个序列随机配对）。
    问题："Q k"——模型读到 k 时要预测它在草堆里对应的 v。
    训练：每个序列插入 2–6 根针，结尾问两倍多个问题（"Q k v Q k' v' …"，从这些针里有放回地抽），只在答案处计损失。
    测试（needle_batch）：要找的针放在草堆的指定深度（0% = 最前面，100% = 紧挨着问题），另有 3 根干扰针，结尾问一次。
    【为什么这样设计】答案只能从针里"抄"——键值配对每个序列都不一样，没法死记；
        这正是 gkamradt 的大海捞针测试（slides）的最小版本，也是"归纳头"（induction head）要做的事：
        找到前面出现过的同一个 k，把它后面的那个 token 抄过来。

【本文件提供】
    NeedleTask（train_batch、needle_batch）
    注意力模式：attention_mask（full / window / strided / dilated），pattern_pairs（每种模式要算多少对）
    LongLM：RMSNorm + RoPE（可调 base、位置插值 scale）/ 无位置编码 + SwiGLU 的小 Transformer（d = 128、2 层、2 个头），可换注意力模式
    train_needle、niah_grid（长度 × 深度的准确率表）、save_model / load_model
    环形注意力：ring_attention（把序列分给 P 块"设备"，K、V 沿环传递，用在线 softmax 合并），RingLedger
    代价公式：flops_per_token（训练时每个 token 的 FLOPs，分成参数部分和注意力部分）、activation_bytes（Ring Attention 论文表 1）

【边界】
    - 稀疏模式用一个 (T, T) 的掩码实现：**数学上**等价于稀疏注意力，但 CPU 上仍按稠密矩阵计算，省不了时间；
      真正省计算要专门的核函数（只算需要的块）。代价用 pattern_pairs 数出来。
    - LongLM 只是一个很小的模型（d = 128、2 层，约 37 万参数），结论是"机制上会发生什么"，不是大模型上的定量结果。
"""

import math
import os
import random
import sys

import torch
import torch.nn as nn
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
for _d in ("lecture18", "lecture04"):
    _p = os.path.join(HERE, "..", _d)
    if _p not in sys.path:
        sys.path.insert(0, _p)

from flash_lab import block_state, finalize, merge_states  # noqa: E402
from modern_lm import RMSNorm, SwiGLU, apply_rope_half, rope_angles  # noqa: E402

CKPT = os.path.join(HERE, "checkpoints")
CURRICULUM_T = 64          # 课程学习第一阶段（没有草堆）的序列长度：放得下 6 根针和 12 个问题（1 + 18 + 36 = 55）


# =====================================================================
# 1. 大海捞针的玩具世界
# =====================================================================

class NeedleTask:
    """
    【参数】n_keys、n_values、n_filler：键、值、填充词的个数；seed：填充词马尔可夫链的种子
    【属性】vocab_size；PAD、BOS、Q、N 和各类 token 的编号区间
    【方法】
        train_batch(gen, B, T, min_needles=2, max_needles=6, hay=True) -> (ids (B, T), targets (B, T))：
            targets 只在"Q k"之后、要预测 v 的位置上是 v，其余为 −100（不计损失）；
            每根针平均被问 q_per_needle 次（有放回地抽）；
            hay=False 时是课程学习的第一阶段：不放草堆（针一根接一根），后面补 <pad>；
            zone=(a, b)：针只插在草堆的 [a, b] 这一段（按比例）。zone=(0.8, 1.0) 得到"序列很长、但依赖都很短"的数据
        needle_batch(gen, B, T, depth, n_distractors=3) -> (ids, answer, distance)：要找的针放在深度 depth（0…1），
            另有 3 根干扰针（位置随机），结尾"Q k"，预测 v
    【边界】T 至少要放得下针和问题（训练时 T ≥ 1 + 6·max_needles + 若干填充词）。
    """

    def __init__(self, n_keys=16, n_values=16, n_filler=32, seed=0):
        self.PAD, self.BOS, self.Q, self.N = 0, 1, 2, 3
        self.key0 = 4
        self.val0 = self.key0 + n_keys
        self.fil0 = self.val0 + n_values
        self.n_keys, self.n_values, self.n_filler = n_keys, n_values, n_filler
        self.vocab_size = self.fil0 + n_filler
        rng = random.Random(seed)
        # 每个填充词后面只会接 4 种填充词（概率随机）：草堆是有规律的"文字"，模型可以学会预测它，但它和问题无关
        self.next_fill = []
        for _ in range(n_filler):
            succ = rng.sample(range(n_filler), 4)
            w = [rng.random() + 0.2 for _ in succ]
            self.next_fill.append((succ, w))

    def names(self):
        """每个编号的可读名字（只用来打印例子）。"""
        out = ["<pad>", "<s>", "Q", "N"] + [f"k{i}" for i in range(self.n_keys)] + [f"v{i}" for i in range(self.n_values)]
        return out + [f"f{i}" for i in range(self.n_filler)]

    def _haystack(self, rng, n):
        f = rng.randrange(self.n_filler)
        out = []
        for _ in range(n):
            out.append(self.fil0 + f)
            succ, w = self.next_fill[f]
            f = rng.choices(succ, w)[0]
        return out

    def train_batch(self, gen_seed, B, T, min_needles=2, max_needles=6, hay=True, q_per_needle=2, zone=(0.0, 1.0)):
        rng = random.Random(gen_seed)
        ids = torch.full((B, T), self.PAD, dtype=torch.long)
        tgt = torch.full((B, T), -100, dtype=torch.long)
        for b in range(B):
            n = rng.randint(min_needles, max_needles)
            keys = rng.sample(range(self.n_keys), n)
            vals = [rng.randrange(self.n_values) for _ in keys]
            nq = q_per_needle * n
            hay_len = T - 1 - 3 * n - 3 * nq if hay else 0
            hay = self._haystack(rng, hay_len)
            lo, hi = round(zone[0] * hay_len), round(zone[1] * hay_len)
            spots = sorted(rng.randint(lo, hi) for _ in range(n))   # 针插在草堆的这些位置之前
            seq, last = [self.BOS], 0
            order = list(range(n))
            rng.shuffle(order)
            for i, s in zip(order, spots):
                seq += hay[last:s] + [self.N, self.key0 + keys[i], self.val0 + vals[i]]
                last = s
            seq += hay[last:]
            # 问题是"有放回地"抽的：如果每根针恰好问一次，模型可以用排除法（最后一个问题只剩一根没问过的针）
            # 答对一半以上的问题，而根本不去比对键——第一版就卡在了这个捷径上
            q_order = [rng.randrange(n) for _ in range(nq)]
            for i in q_order:
                seq += [self.Q, self.key0 + keys[i]]
                tgt[b, len(seq) - 1] = self.val0 + vals[i]       # 读到"Q k"的 k 时，预测下一个 token = v
                seq.append(self.val0 + vals[i])
            ids[b, :len(seq)] = torch.tensor(seq)            # 没有草堆时序列较短，后面留着 <pad>（不计损失）
        return ids, tgt

    def needle_batch(self, gen_seed, B, T, depth, n_distractors=3):
        """
        【参数】depth：要找的那根针在草堆里的相对位置，0 = 紧接在 <s> 后面，1 = 紧挨着结尾的问题；
            n_distractors：另外再插几根"干扰针"（不同的键，位置随机）——只插一根针时，模型不必比对键，
            只要抄"N 后面第二个 token"就能答对，测不出真正的检索能力
        【返回】(ids, answer, distance)：distance 是目标针里的 v 和最后一个位置之间大约隔了多少个 token
        """
        rng = random.Random(gen_seed)
        n = 1 + n_distractors
        ids = torch.empty((B, T), dtype=torch.long)
        ans = torch.empty(B, dtype=torch.long)
        hay_len = T - 1 - 3 * n - 2
        pos = round(depth * hay_len)
        for b in range(B):
            keys = rng.sample(range(self.n_keys), n)
            vals = [rng.randrange(self.n_values) for _ in keys]
            hay = self._haystack(rng, hay_len)
            spots = [pos] + [rng.randint(0, hay_len) for _ in range(n_distractors)]
            order = sorted(range(n), key=lambda i: (spots[i], i))
            seq, last = [self.BOS], 0
            for i in order:
                seq += hay[last:spots[i]] + [self.N, self.key0 + keys[i], self.val0 + vals[i]]
                last = spots[i]
            seq += hay[last:] + [self.Q, self.key0 + keys[0]]
            ids[b] = torch.tensor(seq)
            ans[b] = self.val0 + vals[0]
        distance = T - 1 - (1 + pos + 2)
        return ids, ans, distance


# =====================================================================
# 2. 注意力模式
# =====================================================================

def attention_mask(kind, T, window=16, stride=8, segments=((8, 1), (32, 4), (128, 16))):
    """
    (T, T) 的布尔矩阵：allowed[i, j] = 第 i 个查询能不能看第 j 个键（都满足因果：j ≤ i）。

    【参数】kind：
        "full"    —— 普通因果注意力；
        "window"  —— 滑动窗口（第 4 讲 §6）：只看最近 window 个 token（含自己）；
        "strided" —— Sparse Transformer 的"跨步"模式（Child et al., 2019）：最近 stride 个 token，
                     再加上所有和自己相隔 stride 整数倍的 token；stride ≈ √T 时每个查询看 O(√T) 个键；
        "dilated" —— LongNet 的空洞注意力（Ding et al., 2023）的简化版：把序列切成长度 w 的段，查询只看同一段里
                     每隔 r 个取一个的键；几种 (w, r) 合起来——近处看得密、远处看得疏。
                     （LongNet 原文把查询也按间隔 r 抽稀，再让不同的头用不同的起点，保证每个位置都被覆盖；
                     我们的模型只有 2 个头，所以只抽稀键，让每个查询在每种段长里都能看到一些远处的键。）
    【边界】稀疏模式里每个查询至少能看到自己（对角线永远允许），所以 softmax 不会整行都是 −∞。
    """
    i = torch.arange(T).unsqueeze(1)
    j = torch.arange(T).unsqueeze(0)
    causal = j <= i
    if kind == "full":
        return causal
    if kind == "window":
        return causal & (i - j < window)
    if kind == "strided":
        return causal & ((i - j < stride) | ((i - j) % stride == 0))
    if kind == "dilated":
        allowed = torch.zeros(T, T, dtype=torch.bool)
        for w, r in segments:
            same_seg = (i // w) == (j // w)
            allowed |= same_seg & ((j % w) % r == 0)
        return causal & (allowed | (i == j))
    raise ValueError(kind)


def pattern_pairs(kind, T, **kw):
    """一种注意力模式要计算多少个 (查询, 键) 对——注意力的计算量和显存都正比于它。"""
    return int(attention_mask(kind, T, **kw).sum())


# =====================================================================
# 3. 模型
# =====================================================================

class LongAttention(nn.Module):
    """
    多头因果注意力，位置编码可换：
        pos = "rope"：RoPE（第 4 讲），base 可调；scale > 1 时做**位置插值**（Chen et al., 2023）：
                      把位置 m 换成 m / scale，于是长度 scale·L 的序列被"压"回训练时见过的角度范围 [0, L)；
        pos = "none"：不加位置编码（NoPE）——只靠因果掩码。
    【参数】d_model、n_heads、pos、base、mask_kind（见 attention_mask）、mask_kw
    【属性】scale、base 可以在训练后修改（零样本地换一种位置编码去测更长的序列）
    """

    def __init__(self, d_model, n_heads, pos="rope", base=10000.0, mask_kind="full", mask_kw=None):
        super().__init__()
        self.h, self.dk = n_heads, d_model // n_heads
        self.pos, self.base, self.scale = pos, base, 1.0
        self.mask_kind, self.mask_kw = mask_kind, mask_kw or {}
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)
        self._masks = {}
        self.keep_attn, self.last_attn = False, None

    def mask(self, T, device):
        key = (self.mask_kind, T, repr(sorted(self.mask_kw.items())))
        if key not in self._masks:
            allowed = attention_mask(self.mask_kind, T, **self.mask_kw)
            self._masks[key] = torch.zeros(T, T).masked_fill(~allowed, float("-inf"))
        return self._masks[key].to(device)

    def forward(self, x):
        B, T, D = x.shape
        q, k, v = self.qkv(x).view(B, T, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)
        if self.pos == "rope":
            cos, sin = rope_angles(torch.arange(T) / self.scale, self.dk, self.base)
            q, k = apply_rope_half(q, cos, sin), apply_rope_half(k, cos, sin)
        att = (q @ k.transpose(-2, -1)) / math.sqrt(self.dk) + self.mask(T, x.device)
        att = F.softmax(att, dim=-1)
        if self.keep_attn:
            self.last_attn = att.detach()
        return self.out((att @ v).transpose(1, 2).reshape(B, T, D))


class LongBlock(nn.Module):
    def __init__(self, d_model, n_heads, d_ff, **attn_kw):
        super().__init__()
        self.n1, self.n2 = RMSNorm(d_model), RMSNorm(d_model)
        self.attn = LongAttention(d_model, n_heads, **attn_kw)
        self.ffn = SwiGLU(d_model, d_ff)

    def forward(self, x):
        x = x + self.attn(self.n1(x))
        return x + self.ffn(self.n2(x))


class LongLM(nn.Module):
    """
    小 Transformer 语言模型（Llama 式：RMSNorm、RoPE、SwiGLU、无偏置）。

    【参数】vocab_size、d_model、n_layers、n_heads、d_ff；pos / base / mask_kind / mask_kw 传给每一层的注意力
    【为什么是 d = 128、2 个头】每个头 64 维。先试的 d = 64、4 个头（每头 16 维，和第 2 讲的 GPT 一样）在最简单的
        键值检索上都学不会（准确率停在 0.25）；每头 64 维时 600 步就学会了——比对 16 种键需要足够宽的 QK 空间。
    【返回】forward(ids) -> logits (B, T, V)
    【方法】set_rope(scale=None, base=None)：改所有层的位置插值倍数和 RoPE 的 base（零样本扩展上下文时用）
    """

    def __init__(self, vocab_size, d_model=128, n_layers=2, n_heads=2, d_ff=352, **attn_kw):
        super().__init__()
        self.tok = nn.Embedding(vocab_size, d_model)
        self.blocks = nn.ModuleList(LongBlock(d_model, n_heads, d_ff, **attn_kw) for _ in range(n_layers))
        self.norm = RMSNorm(d_model)
        self.head = nn.Linear(d_model, vocab_size, bias=False)
        for m in self.modules():
            if isinstance(m, (nn.Linear, nn.Embedding)):
                nn.init.normal_(m.weight, 0.0, 0.02)

    def forward(self, ids):
        x = self.tok(ids)
        for blk in self.blocks:
            x = blk(x)
        return self.head(self.norm(x))

    def set_rope(self, scale=None, base=None):
        for blk in self.blocks:
            if scale is not None:
                blk.attn.scale = scale
            if base is not None:
                blk.attn.base = base


def build_model(task, cfg):
    """
    按配置字典构造模型（构造前先设种子：本机上 torch 的默认种子每个进程不同）。
    【默认种子是 1】种子 0 的初始化在本任务上怎样都学不会检索（见 train_needle 的说明），种子 1、2 都能学会。
    """
    torch.manual_seed(cfg.get("seed", 1))
    kw = {"pos": cfg.get("pos", "rope"), "base": cfg.get("base", 10000.0), "mask_kind": cfg.get("mask", "full"),
          "mask_kw": dict(cfg.get("mask_kw", {}))}
    return LongLM(task.vocab_size, cfg.get("d_model", 128), cfg.get("n_layers", 2), cfg.get("n_heads", 2), **kw)


def train_needle(model, task, steps, T, batch=32, lr=1e-3, seed=0, log=None, log_every=0, warmup=100,
                 data_offset=0, curriculum=0, zone=(0.0, 1.0)):
    """
    在大海捞针的训练数据上训练（每一步现做一批新数据；只在答案处计损失）。
    【参数】steps、T（训练长度）、batch、lr（AdamW，线性预热 + 余弦衰减到 10%）、seed（数据的种子）、
        data_offset：数据种子的偏移（继续训练时用不同的数据）；zone：针插在草堆的哪一段（见 NeedleTask.train_batch）
        curriculum：前多少步用"简单版"数据：不放草堆，只有 2–6 根针和问题（序列长 CURRICULUM_T，后面补 <pad>）。
            【为什么】直接在长草堆里训练，模型卡在"随便抄一根针的值"（准确率 ≈ 0.3）上，试过的 3 个种子 1400 步都学不会比对键：
            能产生梯度的只有答案那几个位置，而第二层的注意力一开始均匀地摊在上百个填充词上，"比对键"这条电路得到的信号太弱。
            没有草堆时，种子 1、2 在 400–600 步就学会了（种子 0 在所有设置下都卡住——这类"突然学会"的能力出不出现，
            和随机初始化有很大关系）。先学会、再换成长草堆，准确率就一直保持住了。
    【返回】[(step, 最近 log_every 步的平均损失, 答案的准确率)] 的列表
    """
    opt = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=0.01)
    curve, run_loss, run_acc, run_n = [], 0.0, 0.0, 0
    model.train()
    for step in range(1, steps + 1):
        frac = step / warmup if step <= warmup else 0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * (step - warmup) / max(1, steps - warmup)))
        for g in opt.param_groups:
            g["lr"] = lr * min(1.0, frac)
        warm = step <= curriculum
        ids, tgt = task.train_batch(seed * 10_000_019 + data_offset + step, batch, CURRICULUM_T if warm else T, hay=not warm,
                                    zone=zone)
        logits = model(ids)
        loss = F.cross_entropy(logits.reshape(-1, logits.size(-1)), tgt.reshape(-1), ignore_index=-100)
        opt.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
        mask = tgt != -100
        run_loss += loss.item()
        run_acc += (logits.argmax(-1)[mask] == tgt[mask]).float().mean().item()
        run_n += 1
        if log_every and step % log_every == 0:
            curve.append((step, run_loss / run_n, run_acc / run_n))
            if log:
                log(f"    step {step:5d}  损失 {run_loss / run_n:.3f}  答案准确率 {run_acc / run_n:.3f}")
            run_loss, run_acc, run_n = 0.0, 0.0, 0
    model.eval()
    return curve


@torch.no_grad()
def niah_grid(model, task, lengths, depths, n=64, seed=12345, batch=64):
    """
    大海捞针的准确率表：grid[(T, depth)] = 准确率（n 个序列，每个只有 1 根针）。
    【边界】所有模型用同一批测试序列（种子由 T、depth 决定），比较才公平。
    """
    model.eval()
    grid = {}
    for T in lengths:
        for di, d in enumerate(depths):
            correct = 0
            for s in range(0, n, batch):
                ids, ans, _ = task.needle_batch(seed + 1000 * T + 17 * di + s, min(batch, n - s), T, d)
                correct += int((model(ids)[:, -1].argmax(-1) == ans).sum())
            grid[(T, d)] = correct / n
    return grid


def save_model(model, name, cfg, extra=None):
    os.makedirs(CKPT, exist_ok=True)
    torch.save({"cfg": cfg, "state": model.state_dict(), "extra": extra or {}}, os.path.join(CKPT, f"{name}.pt"))


def load_model(task, name, cfg):
    """读取缓存；不存在或配置不同时返回 None。"""
    path = os.path.join(CKPT, f"{name}.pt")
    if not os.path.exists(path):
        return None, None
    blob = torch.load(path, weights_only=True)
    if blob["cfg"] != cfg:
        return None, None
    model = build_model(task, cfg)
    model.load_state_dict(blob["state"])
    model.eval()
    return model, blob["extra"]


def get_model(task, name, cfg, steps, T, quick=False, log=print, init_from=None, lr=1e-3, data_offset=0,
              setup=None, curriculum=0, zone=(0.0, 1.0), warmup=100):
    """
    有缓存就读取，没有就训练并缓存。init_from 给了一个模型时，从它的权重开始继续训练（扩展上下文）。
    setup(model)：训练前对模型做的修改（例如改 RoPE 的 base 或插值倍数）。
    quick=True 时只训练 20 步、不读写缓存（冒烟测试）。
    """
    full_cfg = dict(cfg, steps=steps, T=T, lr=lr, data_offset=data_offset, curriculum=curriculum, zone=list(zone),
                    warmup=warmup, init=init_from[1] if init_from else None)
    if not quick:
        model, extra = load_model(task, name, full_cfg)
        if model is not None:
            if setup:
                setup(model)
            if log:
                log(f"  （读取缓存的模型：checkpoints/{name}.pt）")
            return model, extra.get("curve", [])
    model = build_model(task, cfg)
    if init_from is not None:
        model.load_state_dict(init_from[0].state_dict())
    if setup:
        setup(model)
    curve = train_needle(model, task, 20 if quick else steps, T, lr=lr, data_offset=data_offset,
                         curriculum=min(curriculum, 10) if quick else curriculum, zone=zone, warmup=warmup,
                         log_every=0 if quick else max(1, steps // 5), log=log)
    if not quick:
        save_model(model, name, full_cfg, {"curve": curve})
        if log:
            log(f"  （已保存到 checkpoints/{name}.pt）")
    return model, curve


# =====================================================================
# 4. 环形注意力
# =====================================================================

class RingLedger:
    """记录环形传递：每一轮每块设备发给下一块的元素个数，以及每块设备每一轮真正算了多少个 (查询, 键) 对。"""

    def __init__(self, P):
        self.P = P
        self.sent = [0] * P
        self.work = []                  # 每一轮一个列表：各设备算的 (q, k) 对数

    def round_time(self):
        """每一轮所有设备要同步（等最慢的一块），总时间 ∝ Σ_轮 max_设备 工作量。"""
        return sum(max(r) for r in self.work)


def ring_attention(Q, K, V, P, causal=True, layout="contiguous", ledger=None):
    """
    环形注意力（Liu et al., 2023）的模拟。

    【思路】把 N 个 token 分给 P 块"设备"，第 r 块设备持有自己那部分 token 的 Q_r、K_r、V_r。
        一共 P 轮：每一轮，每块设备用**手上当前的** K、V 块和自己的 Q_r 算一块注意力，用在线 softmax 的 (m, ℓ, o)
        合并进自己的累加器（第 18 讲）；然后把手上的 K、V 块发给环上的下一块设备、从上一块设备收一块。
        P 轮之后，每块设备都"见过"了全部的 K、V，得到自己那部分 token 的精确输出。
        K、V 传递的顺序恰好就是它们被用到的顺序（slides："和注意力计算里使用 K、V 块的顺序完全相同"），
        所以通信可以和计算重叠：算这一块的同时，下一块已经在路上。
    【参数】Q、K、V：(N, d)；P：设备数（N 要能被 P 整除）；causal：因果掩码；
        layout："contiguous"（第 r 块设备拿第 r 段连续的 token）或 "striped"（token t 给第 t mod P 块设备；
        Brandon et al., 2023 的条带式注意力，让因果掩码下的工作量在设备之间均衡）
    【返回】(O (N, d), ledger)
    """
    N, d = Q.shape
    if N % P:
        raise ValueError("N 必须能被 P 整除")
    ledger = ledger or RingLedger(P)
    if layout == "contiguous":
        idx = [torch.arange(r * N // P, (r + 1) * N // P) for r in range(P)]
    else:
        idx = [torch.arange(r, N, P) for r in range(P)]
    scale = 1.0 / math.sqrt(d)
    acc = [None] * P
    held = list(range(P))                         # 第 r 块设备当前手上拿的是谁的 K、V
    for rnd in range(P):
        work = []
        for r in range(P):
            qi, kj = idx[r], idx[held[r]]
            S = Q[qi] @ K[kj].T * scale
            n_pairs = S.numel()
            if causal:
                allowed = kj.unsqueeze(0) <= qi.unsqueeze(1)
                n_pairs = int(allowed.sum())
                S = S.masked_fill(~allowed, float("-inf"))
            work.append(n_pairs)
            if n_pairs:                           # 整块都在未来：跳过
                st = block_state(S, V[kj])
                acc[r] = st if acc[r] is None else merge_states(acc[r], st)
        ledger.work.append(work)
        if rnd == P - 1:                           # 最后一轮算完就结束，不必再传
            break
        for r in range(P):                         # K、V 沿环传给下一块设备
            ledger.sent[r] += 2 * len(idx[held[r]]) * d
        held = [held[(r - 1) % P] for r in range(P)]
    O = torch.empty_like(Q)
    for r in range(P):
        O[idx[r]] = finalize(acc[r])[0]
    return O, ledger


# =====================================================================
# 5. 代价公式
# =====================================================================

def flops_per_token(n_params, n_layers, d_model, T, causal=True):
    """
    训练时每个 token 的 FLOPs ≈ 6·N_参数 + 注意力部分。
    【注意力部分】每层前向 QKᵀ 和 PV 各 2·T·d 次运算（每个查询对 T 个键），反向是前向的 2 倍：
        12·L·T·d；因果掩码只算一半：6·L·T·d（FlashAttention 跳过未来的块）。
    【返回】(参数部分, 注意力部分)
    """
    attn = (6 if causal else 12) * n_layers * T * d_model
    return 6 * n_params, attn


def activation_bytes(kind, s, h, n_heads, c=4096, b=1):
    """
    一层要保存的最大激活值（字节，BF16），按 Ring Attention 论文表 1：
        vanilla           ：注意力 2·b·n·s²（每个头一个 s×s 的分数矩阵）
        flash（省内存的注意力）：前馈层 8·b·s·h 最大（中间层 4h，BF16 2 字节）
        bpt（注意力 + 前馈都分块）：2·b·s·h
        ring              ：6·b·c·h（每块设备只存 6 个块：Q、当前 K/V、正在接收的 K/V、输出），和 s 无关
    【参数】s：序列长度；h：隐藏维度；n_heads；c：块大小；b：批大小
    """
    if kind == "vanilla":
        return 2 * b * n_heads * s * s
    if kind == "flash":
        return 8 * b * s * h
    if kind == "bpt":
        return 2 * b * s * h
    if kind == "ring":
        return 6 * b * c * h
    raise ValueError(kind)