#!/usr/bin/env python3
"""8-rank GPU byte-exactness test: compose legacy vs v2 (Step A).

Runs the REAL materialize_prefix_and_reuse_current_{kv,index}_page_slots on
every rank with real NCCL collectives and (for v2) real tai-kernel CUDA-IPC
gathers, and asserts the composed dense buffers and locs are byte-identical
between the legacy per-span path and compose_v2.

Run inside the g0033 syh-dev-new container:
  cd /mnt/beegfs/syh/sglang-stable && \
  SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE=1 \
  PYTHONPATH=python:/mnt/beegfs/syh/tai-kernel/python \
  torchrun --nproc-per-node=8 test/manual/test_cp_shared_kv_compose_v2_8rank.py
"""
from __future__ import annotations

import os

import torch
import torch.distributed as dist

from sglang.srt.environ import envs
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout


class _CpGroupShim:
    """Minimal stand-in for the attention CP GroupCoordinator."""

    def __init__(self, group: dist.ProcessGroup) -> None:
        self.device_group = group
        self.world_size = dist.get_world_size(group)
        self.unique_name = "compose_v2_test"
        self.pynccl_comm = None

    def all_gather_into_tensor(self, out: torch.Tensor, t: torch.Tensor) -> None:
        dist.all_gather_into_tensor(out, t, group=self.device_group)

    def all_reduce(self, t: torch.Tensor) -> torch.Tensor:
        dist.all_reduce(t, group=self.device_group)
        return t


def _build_scenario(rank: int, cp_size: int, device: torch.device):
    """bs=4 with mixed prefix/extend lengths, extends crossing page bounds."""
    page_size = 64
    batch_size = 4
    prefix_lens = [640, 1280, 320, 640]
    extend_lens = [95, 130, 64, 200]
    kv_dim = 656  # fp8 layout bytes/token

    from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
        build_in_seq_page_compute_owners,
    )

    g = torch.Generator().manual_seed(20260612)
    logical_rows = []
    current_locs_all = []
    next_page = 1
    prefix_pages_by_req = []
    for req_id, (prefix_len, extend_len) in enumerate(zip(prefix_lens, extend_lens)):
        prefix_pages = prefix_len // page_size
        prefix_pages_by_req.append(prefix_pages)
        req_pages = list(range(next_page, next_page + prefix_pages))
        next_page += prefix_pages
        owners = build_in_seq_page_compute_owners(
            extend_len=extend_len,
            extend_prefix_len=prefix_len,
            page_size=page_size,
            cp_size=cp_size,
        )
        current_pages = []
        for owner in owners:
            # pick a logical page owned by `owner`: owner = (page-1) % cp_size
            candidate = next_page
            while (candidate - 1) % cp_size != int(owner):
                candidate += 1
            current_pages.append(candidate)
            next_page = candidate + 1
        req_pages.extend(current_pages)
        remaining = extend_len
        for page_offset, logical_page in enumerate(current_pages):
            valid = min(page_size, remaining)
            for off in range(valid):
                current_locs_all.append(
                    (req_id, logical_page * page_size + off)
                )
            remaining -= valid
        logical_rows.append(req_pages)

    max_pages = max(len(r) for r in logical_rows)
    logical_pages = torch.zeros((batch_size, max_pages), dtype=torch.int64)
    for req_id, pages in enumerate(logical_rows):
        logical_pages[req_id, : len(pages)] = torch.tensor(pages, dtype=torch.int64)
    logical_pages = logical_pages.to(device)

    # logical locs per request row (token granularity, -1 padded)
    max_tokens = max(
        p + e for p, e in zip(prefix_lens, extend_lens)
    )
    logical_locs = torch.full((batch_size, max_tokens), -1, dtype=torch.int64)
    loc_req_id_rows = []
    for req_id, pages in enumerate(logical_rows):
        total = prefix_lens[req_id] + extend_lens[req_id]
        locs = []
        for token_idx in range(total):
            page = pages[token_idx // page_size]
            locs.append(page * page_size + token_idx % page_size)
        logical_locs[req_id, : len(locs)] = torch.tensor(locs, dtype=torch.int64)
    logical_locs = logical_locs.to(device)

    physical_pages = (next_page // cp_size + 3)
    # Production KV dtype: float8_e4m3fn (caught index_copy_'s missing fp8
    # kernel) — build payloads as bytes, view as fp8.
    kv_cache = torch.zeros(
        (physical_pages * page_size, 1, kv_dim), dtype=torch.uint8, device=device
    )
    layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank)

    # Fill this rank's owned PREFIX pages with deterministic payloads.
    for req_id, pages in enumerate(logical_rows):
        for slot, logical_page in enumerate(pages[: prefix_pages_by_req[req_id]]):
            if (logical_page - 1) % cp_size != rank:
                continue
            phys_page = (logical_page - 1) // cp_size + 1
            payload = (
                torch.arange(page_size * kv_dim, dtype=torch.int64)
                + logical_page * 131
            ).remainder_(251).to(torch.uint8)
            kv_cache[
                phys_page * page_size : (phys_page + 1) * page_size
            ] = payload.view(page_size, 1, kv_dim).to(device)

    # Current rows owned by this rank (writer = page owner in this scenario,
    # so the symm writer list is the storage owner per current page).
    current_page_writer_ranks = []
    for req_pages, prefix_pages in zip(logical_rows, prefix_pages_by_req):
        for page in req_pages[prefix_pages:]:
            current_page_writer_ranks.append((page - 1) % cp_size)
    cur_locs = [
        loc for req_id, loc in current_locs_all
        if (loc // page_size - 1) % cp_size == rank
    ]
    cur_req = [
        req_id for req_id, loc in current_locs_all
        if (loc // page_size - 1) % cp_size == rank
    ]
    current_locs = torch.tensor(cur_locs, dtype=torch.int64, device=device)
    current_req_id = torch.tensor(cur_req, dtype=torch.int64, device=device)
    current_kv = (
        (
            torch.arange(len(cur_locs), dtype=torch.int64).view(-1, 1, 1) * 7
            + torch.arange(kv_dim, dtype=torch.int64).view(1, 1, -1)
            + rank * 31
        )
        .remainder_(249)
        .to(torch.uint8)
        .to(device)
        .view(torch.float8_e4m3fn)
    )
    kv_cache = kv_cache.view(torch.float8_e4m3fn)

    prefix_slot_spans = runtime.build_batch_prefix_slot_spans(
        logical_pages=logical_pages,
        prefix_lens_cpu=prefix_lens,
        page_size=page_size,
    )
    current_slot_spans = runtime.build_batch_current_slot_spans(
        logical_pages=logical_pages,
        prefix_lens_cpu=prefix_lens,
        extend_lens_cpu=extend_lens,
        page_size=page_size,
    )
    slot_remap = runtime.build_shared_token_kv_slot_remap(
        kv_cache,
        logical_locs,
        logical_pages,
        layout,
        page_size,
    )
    loc_req_id = torch.repeat_interleave(
        torch.arange(batch_size, dtype=torch.int64, device=device),
        torch.tensor(
            [max_tokens] * batch_size, dtype=torch.int64, device=device
        ),
    )

    return dict(
        kv_cache=kv_cache,
        logical_pages=logical_pages,
        logical_locs=logical_locs.reshape(-1),
        loc_req_id=loc_req_id,
        current_kv=current_kv,
        current_locs=current_locs,
        current_req_id=current_req_id,
        slot_remap=slot_remap,
        layout=layout,
        page_size=page_size,
        prefix_slot_spans=prefix_slot_spans,
        current_slot_spans=current_slot_spans,
        current_page_writer_ranks=current_page_writer_ranks,
    )


def _build_index_scenario(rank: int, cp_size: int, device: torch.device, s):
    """Index page buffer + current K/scale rows over the bs=4 geometry."""
    page_size = s["page_size"]
    index_page_nbytes = page_size * (128 + 4)
    physical_pages = int(s["kv_cache"].shape[0]) // page_size
    page_buffer = torch.zeros(
        (physical_pages, index_page_nbytes), dtype=torch.uint8, device=device
    )
    # Owned PREFIX pages get deterministic payloads (mirrors the KV pool).
    slot_pages = s["slot_remap"].slot_logical_pages.reshape(-1).cpu()
    prefix_slots = set()
    for start, end in s["prefix_slot_spans"]:
        prefix_slots.update(range(int(start), int(end)))
    for slot in sorted(prefix_slots):
        logical_page = int(slot_pages[slot])
        if logical_page <= 0 or (logical_page - 1) % cp_size != rank:
            continue
        phys_page = (logical_page - 1) // cp_size + 1
        page_buffer[phys_page] = (
            (torch.arange(index_page_nbytes, dtype=torch.int64) + logical_page * 37)
            .remainder_(251)
            .to(torch.uint8)
            .to(device)
        )
    rows = int(s["current_locs"].numel())
    current_k = (
        (
            torch.arange(rows, dtype=torch.int64).view(-1, 1) * 13
            + torch.arange(128, dtype=torch.int64).view(1, -1)
            + rank * 17
        )
        .remainder_(249)
        .to(torch.uint8)
        .to(device)
    )
    current_scale = (
        torch.arange(rows, dtype=torch.float32, device=device).view(-1, 1) * 0.25
        + 1.0
        + rank
    )
    from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as rt

    slot_remap = rt.build_shared_paged_buffer_slot_remap(
        page_buffer,
        s["logical_pages"],
        s["layout"],
    )
    return dict(
        page_buffer=page_buffer,
        current_k=current_k,
        current_scale=current_scale,
        current_locs=s["current_locs"],
        current_req_id=s["current_req_id"],
        slot_remap=slot_remap,
    )


def _check_prefetch_symm(rank: int, cp_size: int, device: torch.device) -> None:
    """bs=1 prefetch-path symm consume vs the legacy sync compose (MLA+index)."""
    from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as rt
    from sglang.srt.layers.attention.nsa.cp_shared_kv_prefetch import (
        CpSharedKVIndexPrefetcher,
        CpSharedKVIndexPrefetchHandle,
        CpSharedKVMlaPrefetcher,
        CpSharedKVMlaPrefetchHandle,
    )
    from sglang.srt.mem_cache.cp_shared_kv_compute_owner import (
        build_in_seq_page_compute_owners,
    )

    page_size = 64
    kv_dim = 656
    index_page_nbytes = page_size * (128 + 4)
    prefix_len, extend_len = 640, 200
    prefix_pages = prefix_len // page_size
    writers = [
        int(o)
        for o in build_in_seq_page_compute_owners(
            extend_len=extend_len,
            extend_prefix_len=prefix_len,
            page_size=page_size,
            cp_size=cp_size,
        )
    ]
    pages = list(range(1, prefix_pages + 1))
    next_page = prefix_pages + 1
    for owner in writers:
        while (next_page - 1) % cp_size != owner:
            next_page += 1
        pages.append(next_page)
        next_page += 1
    total_slots = len(pages)
    logical_pages = torch.tensor([pages], dtype=torch.int64, device=device)
    total_tokens = prefix_len + extend_len
    locs = torch.tensor(
        [
            pages[i // page_size] * page_size + i % page_size
            for i in range(total_tokens)
        ],
        dtype=torch.int64,
        device=device,
    )
    physical_pages = (max(pages) - 1) // cp_size + 3
    layout = CpSharedKVLayout(page_size=page_size, cp_size=cp_size, cp_rank=rank)

    kv_cache = torch.zeros(
        (physical_pages * page_size, 1, kv_dim), dtype=torch.uint8, device=device
    )
    page_buffer = torch.zeros(
        (physical_pages, index_page_nbytes), dtype=torch.uint8, device=device
    )
    for logical_page in pages[:prefix_pages]:
        if (logical_page - 1) % cp_size != rank:
            continue
        phys = (logical_page - 1) // cp_size + 1
        kv_cache[phys * page_size : (phys + 1) * page_size] = (
            (torch.arange(page_size * kv_dim, dtype=torch.int64) + logical_page * 53)
            .remainder_(251)
            .to(torch.uint8)
            .view(page_size, 1, kv_dim)
            .to(device)
        )
        page_buffer[phys] = (
            (torch.arange(index_page_nbytes, dtype=torch.int64) + logical_page * 29)
            .remainder_(251)
            .to(torch.uint8)
            .to(device)
        )

    cur_locs_all = [
        pages[prefix_pages + i // page_size] * page_size + i % page_size
        for i in range(extend_len)
    ]
    mine = [
        i for i in range(extend_len) if writers[i // page_size] == rank
    ]
    current_locs = torch.tensor(
        [cur_locs_all[i] for i in mine], dtype=torch.int64, device=device
    )
    rows = len(mine)
    current_kv = (
        (
            torch.arange(rows, dtype=torch.int64).view(-1, 1, 1) * 11
            + torch.arange(kv_dim, dtype=torch.int64).view(1, 1, -1)
            + rank * 19
        )
        .remainder_(247)
        .to(torch.uint8)
        .to(device)
        .view(torch.float8_e4m3fn)
    )
    kv_cache = kv_cache.view(torch.float8_e4m3fn)
    current_k = (
        (
            torch.arange(rows, dtype=torch.int64).view(-1, 1) * 23
            + torch.arange(128, dtype=torch.int64).view(1, -1)
            + rank * 7
        )
        .remainder_(245)
        .to(torch.uint8)
        .to(device)
    )
    current_scale = (
        torch.arange(rows, dtype=torch.float32, device=device).view(-1, 1) * 0.5
        + 2.0
        + rank
    )
    zeros_req = torch.zeros(rows, dtype=torch.long, device=device)
    loc_req_id = torch.zeros(total_tokens, dtype=torch.long, device=device)

    kv_remap = rt.build_shared_token_kv_slot_remap(
        kv_cache, locs.view(1, -1), logical_pages, layout, page_size
    )
    idx_remap = rt.build_shared_paged_buffer_slot_remap(
        page_buffer, logical_pages, layout
    )
    spans = dict(
        prefix_slot_spans=[(0, prefix_pages)],
        current_slot_spans=[(prefix_pages, total_slots)],
    )

    # --- references via the legacy sync compose (pure NCCL, no barriers) ---
    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
        ref_kv, ref_locs = rt.materialize_prefix_and_reuse_current_kv_page_slots(
            kv_cache=kv_cache,
            logical_locs=locs,
            current_kv_cache=current_kv,
            current_locs=current_locs,
            slot_remap=kv_remap,
            layout=layout,
            page_size=page_size,
            prefix_pages=0,
            loc_req_id=loc_req_id,
            current_req_id=zeros_req,
            layer_id=60,
            **spans,
        )
        ref_idx, ref_pages = rt.materialize_prefix_and_reuse_current_index_page_slots(
            page_buffer=page_buffer,
            current_index_k=current_k,
            current_index_scale=current_scale,
            current_locs=current_locs,
            slot_remap=idx_remap,
            layout=layout,
            page_size=page_size,
            index_head_dim=128,
            prefix_pages=0,
            current_req_id=zeros_req,
            layer_id=60,
            **spans,
        )
    torch.cuda.synchronize()
    dist.barrier()

    # --- prefetchers with a hand-built prefix handle (prefix already
    # materialized + reduced, exactly what start_next_layer_prefix yields) ---
    mla = CpSharedKVMlaPrefetcher(
        layout=layout,
        page_size=page_size,
        prefix_pages=prefix_pages,
        slot_logical_pages=kv_remap.slot_logical_pages,
        page_inverse=kv_remap.page_inverse,
        dense_num_pages=kv_remap.dense_num_pages,
        slot_remap=kv_remap,
        symm_writers=writers,
    )
    dense_kv = kv_cache.new_zeros(
        (kv_remap.dense_num_pages * page_size, 1, kv_dim)
    )
    rt.materialize_local_token_kv_page_slots_into(
        kv_cache=kv_cache,
        dense_kv_cache=dense_kv,
        slot_logical_pages=kv_remap.slot_logical_pages,
        layout=layout,
        page_size=page_size,
        start_slot=0,
        end_slot=prefix_pages,
    )
    prefix_rows = rt.slot_range_to_token_slice(page_size, 0, prefix_pages)
    dist.all_reduce(dense_kv[prefix_rows].view(torch.uint8))
    mla_event = torch.cuda.Event()
    mla_event.record()
    mla.handles[61] = CpSharedKVMlaPrefetchHandle(
        layer_id=61, dense_kv_cache=dense_kv, prefix_rows=prefix_rows,
        event=mla_event,
    )

    idx = CpSharedKVIndexPrefetcher(
        layout=layout,
        prefix_pages=prefix_pages,
        slot_logical_pages=idx_remap.slot_logical_pages,
        page_inverse=idx_remap.page_inverse,
        dense_num_pages=idx_remap.dense_num_pages,
        slot_remap=idx_remap,
        symm_writers=writers,
    )
    dense_pb = page_buffer.new_zeros(
        (idx_remap.dense_num_pages, index_page_nbytes)
    )
    rt.materialize_local_paged_buffer_page_slots_into(
        page_buffer=page_buffer,
        dense_page_buffer=dense_pb,
        slot_logical_pages=idx_remap.slot_logical_pages,
        layout=layout,
        start_slot=0,
        end_slot=prefix_pages,
    )
    dist.all_reduce(dense_pb[1 : prefix_pages + 1])
    idx_event = torch.cuda.Event()
    idx_event.record()
    idx.handles[61] = CpSharedKVIndexPrefetchHandle(
        layer_id=61, dense_page_buffer=dense_pb,
        prefix_rows=slice(1, prefix_pages + 1), event=idx_event,
    )
    torch.cuda.synchronize()
    dist.barrier()

    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
        True
    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
        out_idx = idx.consume_prefix_with_current(
            layer_id=61,
            logical_pages=logical_pages,
            current_index_k=current_k,
            current_index_scale=current_scale,
            current_locs=current_locs,
            page_size=page_size,
            index_head_dim=128,
            current_req_id=zeros_req,
        )
        out_mla = mla.consume_prefix_with_current(
            layer_id=61,
            kv_cache=kv_cache,
            logical_locs=locs,
            current_kv_cache=current_kv,
            current_locs=current_locs,
            loc_req_id=loc_req_id,
            current_req_id=zeros_req,
        )
    torch.cuda.synchronize()
    assert out_idx is not None and out_mla is not None, (
        f"rank{rank}: prefetch consume unexpectedly missed"
    )
    sym_idx, sym_pages = out_idx
    sym_kv, sym_locs = out_mla
    assert torch.equal(ref_pages, sym_pages), f"rank{rank}: prefetch index pages"
    assert torch.equal(ref_locs, sym_locs), f"rank{rank}: prefetch mla locs"
    if not torch.equal(ref_idx, sym_idx):
        bad = (ref_idx != sym_idx).any(dim=-1).nonzero().reshape(-1)[:8]
        raise AssertionError(
            f"rank{rank}: prefetch index symm mismatch at pages {bad.cpu().tolist()}"
        )
    ref_kv_b = ref_kv.view(torch.uint8)
    sym_kv_b = sym_kv.view(torch.uint8)
    if not torch.equal(ref_kv_b, sym_kv_b):
        diff = (ref_kv_b != sym_kv_b).any(dim=-1).any(dim=-1)
        bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
        raise AssertionError(
            f"rank{rank}: prefetch mla symm mismatch at rows {bad} "
            f"(of {int(diff.sum())})"
        )


def _compose(s, layer_id: int, *, writers: list[int] | None = None):
    return runtime.materialize_prefix_and_reuse_current_kv_page_slots(
        kv_cache=s["kv_cache"],
        logical_locs=s["logical_locs"],
        current_kv_cache=s["current_kv"],
        current_locs=s["current_locs"],
        slot_remap=s["slot_remap"],
        layout=s["layout"],
        page_size=s["page_size"],
        prefix_pages=0,
        loc_req_id=s["loc_req_id"],
        current_req_id=s["current_req_id"],
        prefix_slot_spans=s["prefix_slot_spans"],
        current_slot_spans=s["current_slot_spans"],
        layer_id=layer_id,
        current_page_writer_ranks=writers,
    )


def main() -> None:
    dist.init_process_group("nccl")
    rank = dist.get_rank()
    world = dist.get_world_size()
    local_rank = int(os.environ.get("LOCAL_RANK", rank))
    torch.cuda.set_device(local_rank)
    device = torch.device("cuda", local_rank)

    shim = _CpGroupShim(dist.group.WORLD)
    runtime.get_attention_cp_group = lambda: shim  # type: ignore[assignment]

    s = _build_scenario(rank, world, device)

    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
        ref_kv, ref_locs = _compose(s, layer_id=0)
    torch.cuda.synchronize()
    dist.barrier()

    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(True):
        v2_kv, v2_locs = _compose(s, layer_id=0)
    torch.cuda.synchronize()
    dist.barrier()

    assert torch.equal(ref_locs, v2_locs), (
        f"rank{rank}: locs mismatch legacy vs v2"
    )
    ref_kv = ref_kv.view(torch.uint8)
    v2_kv = v2_kv.view(torch.uint8)
    if not torch.equal(ref_kv, v2_kv):
        diff = (ref_kv != v2_kv).any(dim=-1).any(dim=-1)
        bad_rows = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
        raise AssertionError(
            f"rank{rank}: dense kv mismatch at rows {bad_rows} "
            f"(of {int(diff.sum())} differing rows)"
        )

    # Cross-rank: every rank must hold the SAME composed buffer.
    ref_sum = ref_kv.view(torch.uint8).to(torch.float64).sum()
    sums = torch.zeros(world, dtype=torch.float64, device=device)
    dist.all_gather_into_tensor(
        sums, ref_sum.reshape(1).to(device)
    )
    assert torch.allclose(sums, sums[0].expand_as(sums)), (
        f"rank{rank}: composed buffers differ across ranks: {sums.tolist()}"
    )

    if rank == 0:
        print(
            "PASS: legacy vs compose_v2 byte-identical "
            f"(dense rows={int(ref_kv.shape[0])}, world={world})",
            flush=True,
        )

    # ---- Step B: compact symm staging (fill-to-staging + barrier + one
    # mega slot-dense gather,
    # zero NCCL in the current-page phase). Multiple layers exercise the
    # staging parity halves; the arena is exercised in a second pass to
    # prove the two knobs are independent. ----
    def _check_symm(layer_id: int, tag: str) -> None:
        symm_kv, symm_locs = _compose(
            s, layer_id=layer_id, writers=s["current_page_writer_ranks"]
        )
        symm_kv = symm_kv.view(torch.uint8)
        torch.cuda.synchronize()
        assert torch.equal(ref_locs, symm_locs), (
            f"rank{rank} layer{layer_id} [{tag}]: symm locs mismatch"
        )
        if not torch.equal(ref_kv, symm_kv):
            diff = (ref_kv != symm_kv).any(dim=-1).any(dim=-1)
            bad = torch.nonzero(diff).reshape(-1)[:8].cpu().tolist()
            raise AssertionError(
                f"rank{rank} layer{layer_id} [{tag}]: symm dense kv mismatch "
                f"at rows {bad} (of {int(diff.sum())})"
            )

    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
        True
    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
        for layer_id in range(4):
            _check_symm(layer_id, "no-arena")
        with envs.SGLANG_CP_SHARED_KV_COMPOSE_ARENA.override(True):
            for layer_id in range(4, 8):
                _check_symm(layer_id, "arena")
    dist.barrier()
    from sglang.srt.layers.attention.nsa.cp_shared_kv_compose import (
        get_compose_staging,
    )

    staging = get_compose_staging(device)
    assert staging.registered, "symm staging was not registered"
    if rank == 0:
        staged_mb = (
            2
            * staging.capacity_pages
            * (staging.page_nbytes("token_kv") + staging.page_nbytes("index"))
            / (1 << 20)
        )
        print(
            "PASS: compact-staging symm compose byte-identical to v2 across "
            f"8 layers (staging {staging.capacity_pages} pages ~{staged_mb:.1f} "
            "MiB, fill-to-staging + barrier + mega gather engaged; arena on/off)",
            flush=True,
        )

    # ---- Phase 3: index sync compose, legacy vs symm (staging already
    # registered by the token phase — the index compose never registers). ----
    si = _build_index_scenario(rank, world, device, s)

    def _compose_index(layer_id: int, *, writers=None):
        return runtime.materialize_prefix_and_reuse_current_index_page_slots(
            page_buffer=si["page_buffer"],
            current_index_k=si["current_k"],
            current_index_scale=si["current_scale"],
            current_locs=si["current_locs"],
            slot_remap=si["slot_remap"],
            layout=s["layout"],
            page_size=s["page_size"],
            index_head_dim=128,
            prefix_pages=0,
            current_req_id=si["current_req_id"],
            prefix_slot_spans=s["prefix_slot_spans"],
            current_slot_spans=s["current_slot_spans"],
            layer_id=layer_id,
            current_page_writer_ranks=writers,
        )

    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(False):
        ref_idx, ref_dense_pages = _compose_index(40)
    torch.cuda.synchronize()
    dist.barrier()
    with envs.SGLANG_CP_SHARED_KV_COMPOSE_V2.override(
        True
    ), envs.SGLANG_CP_SHARED_KV_COMPOSE_SYMM.override(True):
        for layer_id in range(41, 45):
            symm_idx, symm_dense_pages = _compose_index(
                layer_id, writers=s["current_page_writer_ranks"]
            )
            torch.cuda.synchronize()
            assert torch.equal(ref_dense_pages, symm_dense_pages), (
                f"rank{rank} layer{layer_id}: index dense_pages mismatch"
            )
            if not torch.equal(ref_idx, symm_idx):
                bad = (
                    (ref_idx != symm_idx).any(dim=-1).nonzero().reshape(-1)[:8]
                )
                raise AssertionError(
                    f"rank{rank} layer{layer_id}: index symm mismatch at pages "
                    f"{bad.cpu().tolist()}"
                )
    dist.barrier()
    if rank == 0:
        print("PASS: index sync symm compose byte-identical (4 layers)", flush=True)

    # ---- Phase 4+5: prefetch-path symm consume (MLA + index), bs=1. ----
    _check_prefetch_symm(rank, world, device)
    dist.barrier()
    if rank == 0:
        print(
            "PASS: prefetch consume_prefix_with_current symm (MLA + index) "
            "byte-identical to the sync compose",
            flush=True,
        )
    dist.destroy_process_group()


if __name__ == "__main__":
    main()