"""One decode or speculative-verify step of Muse Spark Open's attention, between the QKV
projection and the output projection, for a batch of requests with a paged FP8 E4M3 KV cache.

Semantics follow transformers' MuseSparkTextAttention; the cache layout follows SGLang
(commit a55ab7d5) with FA3, whose FP8 path casts the normalized query to E4M3.
Inputs are the fused QKV-and-gate projection output and the step's metadata (positions,
cache rows, lengths), which is built outside the timed call.
"""

import math

import torch

from checks import compare, query_precision_bound, relative_l2

# The model's config.
HEAD_DIM = 64
HEADS = {"tp8": (16, 2), "dp": (128, 16)}  # (query heads, key/value heads) on one GPU
SLIDING_WINDOW = 2048
ROPE_THETA = 500_000.0
QK_SCALE_FACTOR = 5.4730064863875
SM_SCALE = HEAD_DIM ** -0.5
RMS_EPS = 1e-5
# Tokens per cache page. Scored cases use 1 (one scattered row per token); some tests use 64.
PAGE = 64
FP8 = torch.float8_e4m3fn
# Per-tensor FP8 cache scales: value = e4m3 * scale. Scored cases use 1; some tests use 1/64
# (kv_scale), so kernels must apply them.
K_SCALE = V_SCALE = 1.0
# Output bound per case (relative L2): FP8_QUERY_FACTOR x the reference's own error with
# FP8 E4M3 queries, plus FP8_QUERY_SLACK, within [FP8_QUERY_FLOOR, REL_L2_TOL]
# (checks.query_precision_bound).
REL_L2_TOL = 0.15
FP8_QUERY_FACTOR = 2.5
FP8_QUERY_FLOOR = 0.03
FP8_QUERY_SLACK = 0.015
# The precision the reference rounds queries to.
QUERY_DTYPE = torch.bfloat16
# Bound (relative L2) on the new FP8 cache entries the step writes.
CACHE_TOL = 8e-2


# dist="soft": each KV head's keys share a unit-RMS direction u: k = sqrt(c2) u + sqrt(1 - c2) z,
# still unit RMS, with c2 uniform per head in SOFT_C2[layer]. Attention spreads over many keys.
SOFT_C2 = {"sliding": (0.3, 0.75), "full": (0.5, 0.62)}


# dist="coherent": soft keys, and for each KV head, keys outside the COHERENT_HEAD
# highest-scoring for its queries get COHERENT_ALPHA x a shared unit-RMS direction added to
# their values.
COHERENT_HEAD = 2048
COHERENT_ALPHA = 1.0


def _coherent_tail(v_cache, k_cache, q, seq_lens, page_table, layer, seed, page, scale):
    """Add the shared direction to the values of each request's low-attention keys."""
    hkv, hq = v_cache.shape[-2], q.shape[-2]
    group = hq // hkv
    tail_gen = torch.Generator(device=v_cache.device).manual_seed(seed + 0xC0DE)
    u = torch.randn((hkv, HEAD_DIM), generator=tail_gen, device=v_cache.device)
    u = u * torch.rsqrt(u.pow(2).mean(-1, keepdim=True))
    qn = _rms(q) * QK_SCALE_FACTOR  # [B, L, Hq, 64], as the reference attends (RoPE on sliding)
    if layer == "sliding":
        pos = seq_lens.long()[:, None] + torch.arange(q.shape[1], device=q.device)
        qn = _rope(qn, pos, ROPE_THETA)
    k_flat = k_cache.view(-1, hkv, HEAD_DIM)
    v_flat = v_cache.view(-1, hkv, HEAD_DIM)
    for b in range(len(seq_lens)):
        seq = int(seq_lens[b])
        first = max(0, seq - SLIDING_WINDOW + 1) if layer == "sliding" else 0
        past = torch.arange(first, seq, device=q.device)
        rows = page_table[b].long()[past // page] * page + past % page
        keys = _dequant(k_flat[rows], scale)  # [S, Hkv, 64]
        # Each key's best logit over its KV head's query heads and the block's tokens.
        logits = torch.einsum("lkgd,skd->kgls", qn[b].view(-1, hkv, group, HEAD_DIM), keys)
        best = logits.flatten(1, 2).amax(1)  # [Hkv, S]
        if best.shape[1] <= COHERENT_HEAD:
            continue
        tail = best < best.topk(COHERENT_HEAD, -1).values[:, -1:]  # [Hkv, S]
        vals = _dequant(v_flat[rows], scale) + COHERENT_ALPHA * tail.T[..., None] * u
        v_flat[rows] = _quant(vals, scale)


# spread="agentic": lengths uniform in ctx +/- AGENTIC_SPREAD (the scored batched cases).
# spread="mixed" (tests): log-normal around ctx, up to 4x it.
AGENTIC_SPREAD = 0.075


def _longest(ctx: int, spread: str) -> int:
    """The longest length _lengths can draw."""
    if spread == "fixed":
        return ctx
    if spread == "agentic":
        return math.ceil(ctx * (1 + AGENTIC_SPREAD))
    return 4 * ctx


def _lengths(gen, batch: int, ctx: int, spread: str, device) -> torch.Tensor:
    if spread == "fixed":
        return torch.full((batch,), ctx, dtype=torch.int64, device=device)
    if spread == "agentic":
        u = torch.rand((batch,), generator=gen, device=device)
        return (ctx * (1 + AGENTIC_SPREAD * (2 * u - 1))).long().clamp(1, _longest(ctx, spread))
    # Log-normal around ctx: many short requests, a long tail.
    z = torch.randn((batch,), generator=gen, device=device)
    return (ctx * torch.exp(z - 0.5)).clamp(16, 4 * ctx).long()


def _shortest_share(spread: str) -> float:
    """A lower bound on mean length / ctx, for the speed-of-light bound."""
    return {"fixed": 1.0, "agentic": 1 - AGENTIC_SPREAD}.get(spread, 0.5)


def _quant(x, scale):
    return (x.float() / scale).to(FP8)


def _dequant(x, scale):
    return x.float() * scale


CACHE_CHUNK = 1 << 20  # cache rows generated at a time


def _fill_cache(gen, shape, device, scale=K_SCALE, seed=0, dist="normal", layer="full"):
    """FP8 key and value caches of `shape`: N(0, 1) entries, soft keys for dist="soft" and
    "coherent" (see SOFT_C2)."""
    hkv = shape[-2]
    k_cache = torch.empty(shape, dtype=FP8, device=device)
    v_cache = torch.empty(shape, dtype=FP8, device=device)
    mix = None
    if dist in ("soft", "coherent"):
        # Its own generator, so the other inputs match the same seed's normal case.
        soft = torch.Generator(device=device).manual_seed(seed + 0x50F7)
        u = torch.randn((hkv, HEAD_DIM), generator=soft, device=device)
        u = u * torch.rsqrt(u.pow(2).mean(-1, keepdim=True))
        lo, hi = SOFT_C2[layer]
        c2 = lo + (hi - lo) * torch.rand((hkv, 1), generator=soft, device=device)
        mix = (c2.sqrt() * u, (1 - c2).sqrt())
    for dst, keys in ((k_cache, True), (v_cache, False)):
        flat = dst.view(-1, hkv, HEAD_DIM)
        for at in range(0, flat.shape[0], CACHE_CHUNK):
            n = min(CACHE_CHUNK, flat.shape[0] - at)
            x = torch.randn((n, hkv, HEAD_DIM), generator=gen, device=device)
            if keys and mix is not None:
                x = mix[0] + mix[1] * x
            flat[at:at + n] = _quant(x, scale)
    return k_cache, v_cache


def _fa3_scheduler_metadata(cache_seqlens, cu_seqlens_q, hq, hkv, page, q_len, layer):
    """FA3's tile schedule for a decode step on a full-attention layer
    (sgl_kernel.flash_attn.get_scheduler_metadata), which SGLang builds once per step from the
    batch's lengths and shares across those layers. None on sliding-window layers and verify
    blocks, where SGLang leaves it to each call."""
    if layer != "full" or q_len != 1:
        return None
    try:
        from sgl_kernel.flash_attn import get_scheduler_metadata
    except ImportError:
        return None
    return get_scheduler_metadata(
        batch_size=cache_seqlens.shape[0], max_seqlen_q=1, max_seqlen_k=int(cache_seqlens.max()),
        num_heads=hq, num_heads_k=hkv, headdim=HEAD_DIM, cache_seqlens=cache_seqlens, qkv_dtype=FP8,
        cu_seqlens_q=cu_seqlens_q, page_size=page, causal=True, num_splits=0)


def generate_input(batch: int, ctx: int, heads: str, layer: str, seed: int, q_len: int = 1,
                   spread: str = "fixed", dist: str = "normal", page: int = PAGE,
                   kv_scale: float = K_SCALE, device: str = "cuda", **_):
    """layer: "sliding" (2,048-token window, RoPE) or "full" (every cached token, no RoPE).
    q_len: 1 for a decode step, 8 for a speculative-decoding verify block (causal inside it).
    dist: "normal", "soft" (see SOFT_C2), "coherent" (see COHERENT_HEAD), or "edges": keys
    aligned with each KV head's queries at the window edges, relative to the block's first
    token (sliding: the oldest position inside the window and the newest just outside it;
    full: position 0)."""
    gen = torch.Generator(device=device).manual_seed(seed)
    hq, hkv = HEADS[heads]
    seq_lens = _lengths(gen, batch, ctx, spread, device)
    pages_per = (seq_lens + q_len + page - 1) // page  # room for the new tokens
    # The pool and page table are sized for the longest possible length, so every
    # repetition's inputs have the same shapes (one CUDA graph is replayed on fresh inputs).
    longest = _longest(ctx, spread)
    max_pages = (longest + q_len + page - 1) // page
    total = batch * max_pages
    # Physical pages are scattered.
    perm = torch.randperm(total, generator=gen, device=device).to(torch.int32)
    page_table = torch.full((batch, max_pages), -1, dtype=torch.int32, device=device)
    offsets = torch.cumsum(pages_per, 0) - pages_per
    for b in range(batch):
        n, o = int(pages_per[b]), int(offsets[b])
        page_table[b, :n] = perm[o:o + n]

    # Cached keys are post-norm (unit RMS per head); values are projections, scale ~1.
    shape = (total, page, hkv, HEAD_DIM)
    k_cache, v_cache = _fill_cache(gen, shape, device, kv_scale, seed, dist, layer)
    q = (2 * torch.randn((batch, q_len, hq, HEAD_DIM), generator=gen, device=device)).to(torch.bfloat16)
    if dist == "coherent":
        _coherent_tail(v_cache, k_cache, q, seq_lens, page_table, layer, seed, page, kv_scale)
    if dist == "edges":
        # Each KV head's key points along its group's first query head of the block's first
        # token, as the kernel will see it (normalized, then RoPE on sliding layers).
        qn = _rms(q[:, 0])
        if layer == "sliding":
            qn = _rope(qn, seq_lens, ROPE_THETA)
        planted = _quant(qn[:, :: hq // hkv], kv_scale)  # [B, Hkv, 64]
        k_flat = k_cache.view(-1, hkv, HEAD_DIM)
        for b in range(batch):
            p = int(seq_lens[b])
            spots = [p - SLIDING_WINDOW + 1, p - SLIDING_WINDOW] if layer == "sliding" else [0]
            for t in spots:
                if 0 <= t < p:
                    k_flat[int(page_table[b, t // page]) * page + t % page] = planted[b]
    k_new = (2 * torch.randn((batch, q_len, hkv, HEAD_DIM), generator=gen, device=device)).to(torch.bfloat16)
    v_new = torch.randn((batch, q_len, hkv, HEAD_DIM), generator=gen, device=device).to(torch.bfloat16)
    gate = torch.randn((batch, q_len, hq * HEAD_DIM), generator=gen, device=device).to(torch.bfloat16)
    # Fused projection output, per token: q | k | v | gate.
    qkv = torch.cat([q.flatten(2), k_new.flatten(2), v_new.flatten(2), gate], -1)
    # Step metadata: each new token's position and flat cache row, per-request cache lengths
    # (including the new tokens) and query offsets, and on full-attention decode steps FA3's
    # tile schedule.
    positions = (seq_lens[:, None] + torch.arange(q_len, device=device)).reshape(-1)
    pos = positions.view(batch, q_len)
    cache_loc = (torch.gather(page_table.long(), 1, pos // page) * page + pos % page).reshape(-1)
    cache_seqlens = (seq_lens + q_len).to(torch.int32)
    cu_seqlens_q = torch.arange(0, batch * q_len + 1, q_len, dtype=torch.int32, device=device)
    return {
        "qkv": qkv,
        "positions": positions,
        "cache_loc": cache_loc,
        "cache_seqlens": cache_seqlens,
        "cu_seqlens_q": cu_seqlens_q,
        "scheduler_metadata": _fa3_scheduler_metadata(cache_seqlens, cu_seqlens_q, hq, hkv, page, q_len, layer),
        "k_cache": k_cache,
        "v_cache": v_cache,
        "k_scale": torch.tensor(kv_scale, dtype=torch.float32, device=device),
        "v_scale": torch.tensor(kv_scale, dtype=torch.float32, device=device),
        "page_table": page_table,
        "seq_lens": seq_lens.to(torch.int32),
        "config": {
            "batch": batch,
            "q_len": q_len,
            "num_q_heads": hq,
            "num_kv_heads": hkv,
            "head_dim": HEAD_DIM,
            "page_size": page,
            "window": SLIDING_WINDOW if layer == "sliding" else 0,
            "rope_theta": ROPE_THETA if layer == "sliding" else 0.0,
            "qk_scale_factor": QK_SCALE_FACTOR,
            "sm_scale": SM_SCALE,
            "kv_scale": kv_scale,
        },
    }


def split_qkv(data):
    """q [B, L, Hq, 64], k and v [B, L, Hkv, 64], gate [B, L, Hq*64] from the fused input."""
    cfg = data["config"]
    hq, hkv, d = cfg["num_q_heads"], cfg["num_kv_heads"], cfg["head_dim"]
    batch, length = cfg["batch"], cfg["q_len"]
    q, k, v, gate = data["qkv"].split([hq * d, hkv * d, hkv * d, hq * d], -1)
    return (q.reshape(batch, length, hq, d), k.reshape(batch, length, hkv, d), v.reshape(batch, length, hkv, d),
            gate)


def _rms(x, eps=RMS_EPS):
    x = x.float()
    return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)


def _rope(x, pos, theta):
    """Interleaved (GPT-J) RoPE: pairs (2i, 2i+1) rotate by pos * theta^(-2i/64).
    x [..., heads, 64] with pos broadcast over everything but the last two dims."""
    inv_freq = 1.0 / (theta ** (torch.arange(0, HEAD_DIM, 2, device=x.device).float() / HEAD_DIM))
    ang = pos.float()[..., None] * inv_freq
    cos, sin = ang.cos().unsqueeze(-2), ang.sin().unsqueeze(-2)
    even, odd = x[..., 0::2], x[..., 1::2]
    return torch.stack([even * cos - odd * sin, even * sin + odd * cos], -1).flatten(-2)


def _rows(data, positions):
    """Flat cache rows of [B, L] positions."""
    size = data["config"]["page_size"]
    page = torch.gather(data["page_table"].long(), 1, positions // size)
    return page * size + positions % size


def new_slots(data) -> torch.Tensor:
    """[B, L] flat cache rows where the block's new tokens go."""
    pos = data["seq_lens"].long()[:, None] + torch.arange(data["config"]["q_len"], device=data["seq_lens"].device)
    return _rows(data, pos)


def ref_kernel(data, query_dtype=None, new_kv=None):
    """Returns (output [B, L, Hq*64], new keys and values [B, L, Hkv, 64] as stored: FP8).
    query_dtype: the precision queries are rounded to (default QUERY_DTYPE). With new_kv
    (FP8 keys and values [B, L, Hkv, 64]), the block attends those as its own entries."""
    cfg = data["config"]
    batch, length = cfg["batch"], cfg["q_len"]
    hq, hkv = cfg["num_q_heads"], cfg["num_kv_heads"]
    q_in, k_in, v_in, gate = split_qkv(data)
    pos = data["seq_lens"].long()[:, None] + torch.arange(length, device=q_in.device)
    q = _rms(q_in) * cfg["qk_scale_factor"]
    k = _rms(k_in)
    if cfg["rope_theta"]:
        q, k = _rope(q, pos, cfg["rope_theta"]), _rope(k, pos, cfg["rope_theta"])
    # FP8 rounds the normalized query, before qk_scale_factor.
    query_dtype = query_dtype or QUERY_DTYPE
    if query_dtype == torch.bfloat16:
        q = q.to(torch.bfloat16).float()
    else:
        q = (q / cfg["qk_scale_factor"]).to(query_dtype).float() * cfg["qk_scale_factor"]
    k_new, v_new = _quant(k, data["k_scale"]), _quant(v_in, data["v_scale"])
    if new_kv is not None:
        k_new, v_new = new_kv
    k_flat = data["k_cache"].view(-1, hkv, HEAD_DIM)
    v_flat = data["v_cache"].view(-1, hkv, HEAD_DIM)
    group = hq // hkv

    out = torch.empty((batch, length, hq, HEAD_DIM), device=q.device)
    for b in range(batch):
        seq = int(data["seq_lens"][b])
        past = torch.arange(seq, device=q.device)
        size = cfg["page_size"]
        past_rows = data["page_table"][b].long()[past // size] * size + past % size
        keys_past = _dequant(k_flat[past_rows], data["k_scale"])
        vals_past = _dequant(v_flat[past_rows], data["v_scale"])
        keys = torch.cat([keys_past, _dequant(k_new[b], data["k_scale"])])  # positions 0 .. seq+L-1
        vals = torch.cat([vals_past, _dequant(v_new[b], data["v_scale"])])
        for j in range(length):
            p = seq + j
            start = max(0, p - cfg["window"] + 1) if cfg["window"] else 0
            kb, vb = keys[start:p + 1], vals[start:p + 1]
            att = torch.einsum("kgd,tkd->kgt", q[b, j].view(hkv, group, HEAD_DIM), kb) * cfg["sm_scale"]
            out[b, j] = torch.einsum("kgt,tkd->kgd", att.softmax(-1), vb).reshape(hq, HEAD_DIM)
    out = _rms(out) * torch.sigmoid(gate.float()).view(batch, length, hq, HEAD_DIM)
    return out.reshape(batch, length, hq * HEAD_DIM), k_new, v_new


def collect(data, output) -> torch.Tensor:
    """Runs in the submission's process after a step: the output followed by the new cache
    entries, as raw bytes (see unpack)."""
    cfg = data["config"]
    shape = (cfg["batch"], cfg["q_len"], cfg["num_q_heads"] * HEAD_DIM)
    if type(output) is not torch.Tensor or output.dtype != torch.bfloat16 or tuple(output.shape) != shape:
        raise ValueError(f"expected a bfloat16 tensor of shape {shape}")
    slots = new_slots(data)
    hkv = cfg["num_kv_heads"]
    k = data["k_cache"].view(-1, hkv, HEAD_DIM)[slots]
    v = data["v_cache"].view(-1, hkv, HEAD_DIM)[slots]
    return torch.cat([p.contiguous().view(torch.uint8).reshape(-1) for p in (output, k, v)])


def unpack(raw: torch.Tensor, cfg: dict):
    b, length, hq, hkv = cfg["batch"], cfg["q_len"], cfg["num_q_heads"], cfg["num_kv_heads"]
    shapes = [((b, length, hq * HEAD_DIM), torch.bfloat16), ((b, length, hkv, HEAD_DIM), FP8),
              ((b, length, hkv, HEAD_DIM), FP8)]
    sizes = [math.prod(shape) * dtype.itemsize for shape, dtype in shapes]
    if raw.dtype != torch.uint8 or raw.numel() != sum(sizes):
        return None
    out, at = [], 0
    for (shape, dtype), n in zip(shapes, sizes):
        out.append(raw[at:at + n].view(dtype).view(shape))
        at += n
    return out


def check_implementation(data, raw):
    parts = unpack(raw, data["config"])
    if parts is None:
        return False, "malformed output"
    out, k, v = parts
    _, ref_k, ref_v = ref_kernel(data)
    for name, got, want, scale in (("key", k, ref_k, data["k_scale"]), ("value", v, ref_v, data["v_scale"])):
        ok, kv_msg = compare(_dequant(got, scale), _dequant(want, scale), CACHE_TOL)
        if not ok:
            return False, f"new {name} cache entries: {kv_msg}"
    # The output is checked against the reference attending the submission's own new cache
    # entries (checked above), since FP8 rounding of those entries varies by implementation.
    ref = ref_kernel(data, new_kv=(k, v))[0].to(torch.bfloat16)
    e4m3 = relative_l2(ref_kernel(data, torch.float8_e4m3fn, new_kv=(k, v))[0].to(torch.bfloat16), ref)
    tol = query_precision_bound(e4m3, FP8_QUERY_FACTOR, FP8_QUERY_FLOOR, REL_L2_TOL, FP8_QUERY_SLACK)
    ok, msg = compare(out, ref, tol)
    msg = (f"{msg} [bound {tol:.3f}: {FP8_QUERY_FACTOR:g} x the reference's {e4m3:.3f} with FP8 E4M3 queries"
           f" + {FP8_QUERY_SLACK:g}]")
    return ok, msg


def speed_of_light_ns(batch: int, ctx: int, heads: str, layer: str, hw: dict, q_len: int = 1,
                      spread: str = "fixed", **_) -> float:
    hq, hkv = HEADS[heads]
    attended = min(ctx, SLIDING_WINDOW) if layer == "sliding" else ctx
    # Uneven lengths average about ctx but can sit below it; stay a lower bound.
    tokens = batch * attended * _shortest_share(spread)
    bytes_moved = tokens * hkv * HEAD_DIM * 2 + batch * q_len * (2 * hq + 2 * hkv) * HEAD_DIM * 2 * 2
    flops = 4.0 * tokens * q_len * hq * HEAD_DIM
    return 1e9 * max(bytes_moved / hw["hbm_bytes_per_s"], flops / hw["bf16_dense_flops_per_s"])