Restructure bs>1 CP shared-KV compose to one gather + one collective

The bs>1 partial-current materialize issued one sum-all-reduce per request span per buffer per layer (48 collectives per F-layer at bs=12, 2880 per batch, 419ms + 49ms launch gaps in the production trace), all inline on the compute stream. The data is a partitioned gather, not a reduction: every byte has exactly one producer. compose_v2 (SGLANG_CP_SHARED_KV_COMPOSE_V2, default on) replaces this with: - fast path: one tai-kernel CUDA-IPC slot-dense gather covering ALL prefix spans (full-range descriptors, -1 sentinels zero-fill current slots and replace the dense zero-fill) + ONE collective over the compact current pages (uint8 byte view; exact because every byte is writer-exclusive). - fallback (no peer IPC): local materialize of all prefix spans + ONE whole-buffer sum-all-reduce (rows are still writer-exclusive pre-reduce). IPC capability is decided once by the cached peer-pointer probe; after a successful probe a failing gather raises (no per-call try/except). cp_shared_kv_compose.py adds the per-batch ComposePlan descriptor cache (layer-invariant, keyed on the slot_logical_pages identity) and the CpComposeArena with tier-S carve discipline (deterministic bump, layer- parity halves; default off) so the Step B symmetric-memory conversion is a registration flip. Microbenchmark (g0034 8xH200, traced 12-req batch, per batch): per-span 214ms -> fused AR 119ms -> IPC prefix + compact current 84ms; symm target 61ms. Validation: 143 unit tests incl. v2-contract twins, legacy siblings and rank-merged simulations under both paths; mem_cache dir 432 passed; 8-rank GPU byte-exactness vs legacy with real NCCL + IPC (test/manual/ test_cp_shared_kv_compose_v2_8rank.py) passed with no fallback markers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
2026-06-11 19:16:28 +00:00
parent ecf7998c80
commit f47b739e30
6 changed files with 1357 additions and 4 deletions
@@ -0,0 +1,270 @@
+#!/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 g0034 cjy-glm5-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.layers.attention.nsa.cp_shared_kv_compose import (
+    reset_compose_plan_cache,
+)
+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)
+    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).
+    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)
+    )
+
+    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_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,
+    )
+
+
+def _compose(s, layer_id: int):
+    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,
+    )
+
+
+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()
+
+    reset_compose_plan_cache()
+    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"
+    )
+    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.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,
+        )
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()