# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Single-process test for KVCacheManagerV2 + KvCacheTransceiverV2.

Uses threading + ThreadSafeDistributed to create KvCacheTransceiverV2 instances
in a single process. Validates KV cache transfer correctness across different
TP/PP/DP/MLA/sliding-window configurations for V2 cache managers.
"""

import gc
import os
import threading
import uuid
from types import SimpleNamespace

# Do not inherit a NIC pin from the host: the selected interface may not exist
# in the test container and would prevent the NIXL agent from initializing.
os.environ.pop("UCX_NET_DEVICES", None)
# Exclude UCX IB transport (avoid NIXL setup hangs without IB) and gdr_copy
# (avoid SIGSEGV at process exit from UCX rcache cleanup; gdr_copy disabled
# falls back to cuda_ipc / cuda_copy without affecting correctness).
# Force a deterministic UCX/NIXL config regardless of what the cluster/CI
# injects; see test_kv_transfer.py for the full rationale.
os.environ["UCX_TLS"] = "^ib,gdr_copy"
os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1"
from dataclasses import dataclass, field
from typing import Dict, List, Optional

import pytest
import torch

import tensorrt_llm
import tensorrt_llm.bindings
import tensorrt_llm.tensorrt_llm_transfer_agent_binding  # noqa: F401
from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams
from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1
from tensorrt_llm._torch.disaggregation.resource.utils import get_global_layer_ids
from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2
from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import (
    KVCacheLayerOwner,
    KVCacheLayerSpec,
    KVCacheManagerV2,
    Role,
)
from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestType
from tensorrt_llm._utils import TensorWrapper, convert_to_torch_tensor, get_size_in_bytes
from tensorrt_llm.bindings import DataType
from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp
from tensorrt_llm.llmapi.llm_args import (
    BlockReuseConfig,
    CacheTransceiverConfig,
    KvPoolRebalanceConfig,
)

AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
NUM_LAYERS = 4
NUM_KV_HEADS = 4  # 1 for MLA
HEAD_DIM = 128
TOKENS_PER_BLOCK = 8
MAX_SEQ_LEN = 256
MAX_BATCH_SIZE = 4
VOCAB_SIZE = 32000
REQUEST_LENGTHS = [30, 60, 80]


def test_assert_disagg_history_declared_passes_when_contract_met():
    req = SimpleNamespace(py_request_id=7, prompt_len=17)
    cache_manager = SimpleNamespace(get_history_length=lambda r: 17)
    transceiver = object.__new__(KvCacheTransceiverV2)
    transceiver._kv_cache_manager = cache_manager

    # Should not raise — history_length matches prompt_len exactly.
    transceiver._assert_disagg_history_declared(req)

    # And also passes when history_length exceeds prompt_len (e.g., already advanced).
    cache_manager.get_history_length = lambda r: 100
    transceiver._assert_disagg_history_declared(req)


def test_assert_disagg_history_declared_raises_on_contract_violation():
    req = SimpleNamespace(py_request_id=7, prompt_len=17)
    cache_manager = SimpleNamespace(get_history_length=lambda r: 0)
    transceiver = object.__new__(KvCacheTransceiverV2)
    transceiver._kv_cache_manager = cache_manager

    with pytest.raises(RuntimeError, match="history_length=0 < prompt_len=17"):
        transceiver._assert_disagg_history_declared(req)


def test_assert_disagg_history_declared_noops_when_cache_released():
    """If the cache was released (e.g., cancelled mid-transfer), skip the check."""
    req = SimpleNamespace(py_request_id=7, prompt_len=17)
    cache_manager = SimpleNamespace(get_history_length=lambda r: None)
    transceiver = object.__new__(KvCacheTransceiverV2)
    transceiver._kv_cache_manager = cache_manager

    transceiver._assert_disagg_history_declared(req)


# ---------------------------------------------------------------------------
# PP layer distribution helpers (mirrors C++ getLayerNumPPRank)
# ---------------------------------------------------------------------------
def _get_layers_per_pp(num_layers: int, pp_size: int) -> List[int]:
    """Return a list of layer counts per PP rank.

    When num_layers is not evenly divisible by pp_size, the first
    (num_layers % pp_size) ranks get one extra layer.
    Matches Mapping.pp_layers / torch.tensor_split behaviour.
    """
    base = num_layers // pp_size
    extra = num_layers % pp_size
    return [base + (1 if r < extra else 0) for r in range(pp_size)]


def _pp_rank_of_layer(layer_idx: int, num_layers: int, pp_size: int) -> int:
    """Return the PP rank that owns *layer_idx*."""
    cumulative = 0
    for pp_rank, count in enumerate(_get_layers_per_pp(num_layers, pp_size)):
        cumulative += count
        if layer_idx < cumulative:
            return pp_rank
    raise ValueError(f"layer_idx {layer_idx} out of range for num_layers={num_layers}")


def _pp_layer_start(pp_rank: int, num_layers: int, pp_size: int) -> int:
    """Return the global layer index of the first layer on *pp_rank*."""
    start = 0
    for r, count in enumerate(_get_layers_per_pp(num_layers, pp_size)):
        if r == pp_rank:
            return start
        start += count
    raise ValueError(f"pp_rank {pp_rank} out of range for pp_size={pp_size}")


@dataclass
class KvCacheConfigV2:
    """KvCacheConfig wrapper with max_util_for_resume for KVCacheManagerV2."""

    max_tokens: Optional[int] = None
    enable_block_reuse: bool = False
    max_attention_window: Optional[List[int]] = None
    sink_token_length: Optional[int] = None
    free_gpu_memory_fraction: Optional[float] = None
    host_cache_size: Optional[int] = None
    disk_cache_size: Optional[int] = None
    disk_cache_path: Optional[str] = None
    onboard_blocks: bool = True
    cross_kv_cache_fraction: Optional[float] = None
    cold_tier_offload_min_priority: Optional[int] = None
    event_buffer_max_size: int = 0
    kv_cache_event_hash_algo: str = "auto"
    max_gpu_total_bytes: Optional[int] = None
    enable_partial_reuse: bool = False
    copy_on_partial_reuse: bool = False
    dtype: str = "auto"
    disk_prefetch_num_reqs: int = 4
    pool_ratio: Optional[List[float]] = None
    avg_seq_len: Optional[int] = None
    block_reuse_config: BlockReuseConfig = field(default_factory=BlockReuseConfig)
    enable_swa_scratch_reuse: bool = False
    max_util_for_resume: float = 0.95
    # Mirrors KvCacheConfig.kv_pool_rebalance_config in
    # tensorrt_llm/llmapi/llm_args.py; KVCacheManagerV2._build_base_config()
    # reads it unconditionally.
    kv_pool_rebalance_config: KvPoolRebalanceConfig = field(default_factory=KvPoolRebalanceConfig)
    enable_kv_pool_rebalance: bool = False


# ---------------------------------------------------------------------------
# ThreadSafeDistributed: threading.Barrier-based Distributed mock
# ---------------------------------------------------------------------------
class ThreadSafeDistributed:
    """Distributed mock using threading.Barrier for single-process multi-rank testing."""

    def __init__(
        self,
        local_rank: int,
        world_size: int,
        tp_size: int,
        pp_size: int,
        tp_rank: int,
        pp_rank: int,
        shared: dict,
        cp_rank: int = 0,
        cp_size: int = 1,
    ):
        self.rank = local_rank
        self._world_size = world_size
        self._tp_size = tp_size
        self._pp_size = pp_size
        self._tp_rank = tp_rank
        self._pp_rank = pp_rank
        # CP is orthogonal to TP/PP: each cp slice runs its own TP/PP collectives,
        # so the per-group barriers are indexed by (rank, cp_rank). ctx side is
        # cp_size==1, preserving the original single-axis behaviour.
        # HELIX VALIDATE: assumes the transceiver's setup collectives group TP/PP
        # within a fixed cp_rank (cp orthogonal). Confirm against the real helix
        # registration exchange on GPU.
        self._cp_rank = cp_rank
        self._cp_size = cp_size
        self._s = shared
        self._bcast_idx = 0
        self._ag_idx = 0
        self._pp_ag_idx = 0
        self._tp_ag_idx = 0

    @property
    def tp_size(self):
        return self._tp_size

    @property
    def pp_size(self):
        return self._pp_size

    @property
    def world_size(self):
        return self._world_size

    def broadcast(self, obj, root=0):
        idx = self._bcast_idx
        self._bcast_idx += 1
        key = f"bcast_{idx}"
        if self.rank == root:
            self._s[key] = obj
        self._s["barrier"].wait()
        result = self._s[key]
        self._s["barrier"].wait()
        return result

    def allgather(self, obj):
        idx = self._ag_idx
        self._ag_idx += 1
        key = f"ag_{idx}"
        with self._s["lock"]:
            if key not in self._s:
                self._s[key] = [None] * self._world_size
            self._s[key][self.rank] = obj
        self._s["barrier"].wait()
        result = list(self._s[key])
        self._s["barrier"].wait()
        return result

    def pp_allgather(self, obj):
        idx = self._pp_ag_idx
        self._pp_ag_idx += 1
        key = f"pp_ag_{idx}_tp{self._tp_rank}_cp{self._cp_rank}"
        with self._s["lock"]:
            if key not in self._s:
                self._s[key] = [None] * self._pp_size
            self._s[key][self._pp_rank] = obj
        # Sync only the PP group that shares this (tp_rank, cp_rank). With attention data
        # parallelism each tp_rank is an independent instance that may run a different number
        # of collectives, so a single global barrier would deadlock; a per-group one does not.
        # CP adds an orthogonal axis, so the group is keyed by cp_rank too.
        pp_barrier = self._s["pp_barriers"][self._tp_rank * self._cp_size + self._cp_rank]
        pp_barrier.wait()
        result = list(self._s[key])
        pp_barrier.wait()
        return result

    def tp_allgather(self, obj):
        idx = self._tp_ag_idx
        self._tp_ag_idx += 1
        key = f"tp_ag_{idx}_pp{self._pp_rank}_cp{self._cp_rank}"
        with self._s["lock"]:
            if key not in self._s:
                self._s[key] = [None] * self._tp_size
            self._s[key][self._tp_rank] = obj
        # Sync only the TP group that shares this (pp_rank, cp_rank) (see pp_allgather).
        tp_barrier = self._s["tp_barriers"][self._pp_rank * self._cp_size + self._cp_rank]
        tp_barrier.wait()
        result = list(self._s[key])
        tp_barrier.wait()
        return result


# ---------------------------------------------------------------------------
# Threading helpers
# ---------------------------------------------------------------------------
def run_concurrent(items, fn):
    """Run fn(item) for each item concurrently in threads and propagate errors."""
    errors = [None] * len(items)
    results = [None] * len(items)

    def _worker(idx, item):
        try:
            results[idx] = fn(item)
        except Exception as e:
            errors[idx] = e

    threads = [threading.Thread(target=_worker, args=(i, item)) for i, item in enumerate(items)]
    for t in threads:
        t.start()
    for t in threads:
        t.join()
    for i, err in enumerate(errors):
        if err is not None:
            raise err
    return results


# ---------------------------------------------------------------------------
# Cache Manager Creation
# ---------------------------------------------------------------------------
def _get_ctx_seed(
    tp_rank: int, pp_rank: int, tp_size: int, num_kv_heads: int, seed_base: int = 1000
) -> int:
    """Compute random seed for ctx pool initialization.

    When num_kv_heads >= tp_size, heads are split across TP ranks -> unique seed per rank.
    When num_kv_heads < tp_size, heads are replicated -> same seed for all TP ranks in a PP group.
    """
    if num_kv_heads >= tp_size:
        return seed_base + pp_rank * tp_size + tp_rank
    else:
        return seed_base + pp_rank


def _create_cache_manager(
    mapping: Mapping,
    is_mla: bool,
    use_v2: bool,
    max_attention_window_vec: Optional[List[int]] = None,
    num_layers: int = NUM_LAYERS,
    max_batch_size: int = MAX_BATCH_SIZE,
    num_draft_layers: int = 0,
) -> "KVCacheManagerV2":
    """Create a KVCacheManagerV2 for the given mapping.

    ``num_draft_layers`` K-only draft-owned layers follow the target's (MLA, pp 1),
    as a standalone MLA drafter's unified draft KV cache lays them out.
    """
    num_kv_heads = 1 if is_mla else NUM_KV_HEADS
    cache_type = CacheTypeCpp.SELFKONLY if is_mla else CacheTypeCpp.SELF

    if max_attention_window_vec is None:
        max_attention_window_vec = [MAX_SEQ_LEN]
    layer_specs = None
    if num_draft_layers:
        assert is_mla and mapping.pp_size == 1
        layer_specs = tuple(
            KVCacheLayerSpec(
                global_layer_id=layer_id,
                owner=(
                    KVCacheLayerOwner.TARGET if layer_id < num_layers else KVCacheLayerOwner.DRAFT
                ),
                dtype=DataType.FLOAT,
                num_kv_heads=1,
                head_dim=HEAD_DIM,
                buffer_roles=(Role.KEY,),
            )
            for layer_id in range(num_layers + num_draft_layers)
        )

    return KVCacheManagerV2(
        KvCacheConfigV2(
            max_tokens=2048,
            enable_block_reuse=False,
            max_attention_window=max_attention_window_vec,
        ),
        cache_type,
        num_layers=num_layers,
        num_kv_heads=num_kv_heads,
        head_dim=HEAD_DIM,
        tokens_per_block=TOKENS_PER_BLOCK,
        max_seq_len=MAX_SEQ_LEN,
        max_batch_size=max_batch_size,
        mapping=mapping,
        dtype=DataType.FLOAT,
        vocab_size=VOCAB_SIZE,
        is_disagg=True,
        layer_specs=layer_specs,
    )


def _create_managers_for_instance(
    tp: int,
    pp: int,
    enable_dp: bool,
    is_mla: bool,
    use_v2: bool,
    max_attention_window_vec: Optional[List[int]] = None,
    num_layers: int = NUM_LAYERS,
    max_batch_size: int = MAX_BATCH_SIZE,
    cp_size: int = 1,
    num_draft_layers: int = 0,
) -> list[KVCacheManagerV2]:
    """Create cache managers for all ranks in an instance.

    cp_size > 1 builds a helix-CP instance: each cp rank owns the global blocks
    strided by [cp_rank::cp_size] (see transfer.py). Managers size off the full
    (global) prompt via total_input_len_cp; the ranks are enumerated
    pp_rank * (tp*cp) + tp_rank * cp + cp_rank.
    """
    managers = []
    cp_config = {"cp_type": "HELIX"} if cp_size > 1 else None
    for pp_rank in range(pp):
        for tp_rank in range(tp):
            for cp_rank in range(cp_size):
                rank = (pp_rank * tp + tp_rank) * cp_size + cp_rank
                mapping = Mapping(
                    world_size=tp * pp * cp_size,
                    rank=rank,
                    tp_size=tp,
                    pp_size=pp,
                    cp_size=cp_size,
                    cp_config=cp_config,
                    enable_attention_dp=enable_dp,
                )
                managers.append(
                    _create_cache_manager(
                        mapping,
                        is_mla,
                        use_v2,
                        max_attention_window_vec,
                        num_layers,
                        max_batch_size,
                        num_draft_layers,
                    )
                )
    return managers


# ---------------------------------------------------------------------------
# Pool Data Initialization
# ---------------------------------------------------------------------------


def _init_pool_data_v2(
    managers: List[KVCacheManagerV2],
    tp: int,
    is_mla: bool,
    fill_random: bool = True,
    seed_base: int = 1000,
):
    """Initialize pool data for V2 managers."""
    num_kv_heads = 1 if is_mla else NUM_KV_HEADS
    for rank, mgr in enumerate(managers):
        pp_rank = rank // tp
        tp_rank = rank % tp
        page_table = KVRegionExtractorV1(mgr).page_table

        unique_pools: Dict[int, int] = {}
        for pg in page_table.pool_groups:
            for pool_desc in pg.pools:
                key = pool_desc.base_address
                pool_bytes = pool_desc.slot_bytes * pool_desc.num_slots
                if key not in unique_pools or pool_bytes > unique_pools[key]:
                    unique_pools[key] = pool_bytes

        element_bytes = get_size_in_bytes(1, mgr.dtype)
        for pool_base_ptr, pool_size in unique_pools.items():
            pool_size_elements = pool_size // element_bytes
            pool_tensor = convert_to_torch_tensor(
                TensorWrapper(pool_base_ptr, mgr.dtype, [pool_size_elements])
            )
            if fill_random:
                seed = _get_ctx_seed(tp_rank, pp_rank, tp, num_kv_heads, seed_base)
                generator = torch.Generator(device=pool_tensor.device).manual_seed(seed)
                random_values = torch.rand(
                    pool_tensor.shape,
                    dtype=pool_tensor.dtype,
                    device=pool_tensor.device,
                    generator=generator,
                )
                pool_tensor.copy_(random_values)
            else:
                pool_tensor.zero_()


def _init_pool_data(managers, tp, is_mla, use_v2, fill_random=True, seed_base=1000):
    """Initialize V2 pool data for every rank."""
    _init_pool_data_v2(managers, tp, is_mla, fill_random, seed_base)


# ---------------------------------------------------------------------------
# Add sequence to manager
# ---------------------------------------------------------------------------
def _make_gen_request(
    gen_rid, req_len, unique_rid, ctx_rid, ctx_dp_rank, ctx_info_endpoint, sampling_params
):
    """Build one GENERATION_ONLY request wired to its context peer.

    Used to mint a distinct request per helix CP rank (each needs its own
    rank-local prompt_len); mirrors the inline construction in run_transfer_test.
    """
    req = LlmRequest(
        request_id=gen_rid,
        max_new_tokens=1,
        input_tokens=list(range(req_len)),
        sampling_config=tensorrt_llm.bindings.SamplingConfig(
            sampling_params._get_sampling_config()
        ),
        is_streaming=False,
        llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY,
    )
    req.py_disaggregated_params = DisaggregatedParams(
        ctx_request_id=ctx_rid,
        ctx_dp_rank=ctx_dp_rank,
        ctx_info_endpoint=ctx_info_endpoint,
        disagg_request_id=unique_rid,
    )
    return req


def _helix_local_len(global_len: int, cp_rank: int, cp_size: int) -> int:
    """Rank-local token count under helix (mirror of KVCacheManagerV2._helix_local_len).

    A ledger block spans tpb*cp_size global tokens; global block b (page ``b``
    within its ledger block) lives on CP rank ``b %% cp_size``. This is the
    length the *transceiver* uses to bound the
    transferred block list -- distinct from the ledger *capacity*, which the
    manager sizes off the global length (rounded up to whole ledger blocks) and
    so may bind one extra trailing page on ranks that own no token in it.
    """
    tpb = TOKENS_PER_BLOCK
    ledger_tpb = tpb * cp_size
    full, rem = divmod(global_len, ledger_tpb)
    return full * tpb + min(max(rem - cp_rank * tpb, 0), tpb)


def _add_sequence(
    mgr,
    request_id: int,
    prompt_len: int,
    use_v2: bool,
    *,
    is_generation: bool = False,
    scratch_tokens: int = 0,
):
    """Add a sequence to the cache manager. Returns kv_cache for V2 (needed for cleanup).

    scratch_tokens > 0 (generation side only) sizes the cache to
    prompt_len + scratch_tokens with history_length still prompt_len -- exactly
    what prepare_disagg_gen_init does for MTP -- so the speculative tail blocks
    are allocated *and their pages bound* (materialised) without a forward pass.
    That reproduces production's "surplus block" on a windowed slice, which the
    harness cannot (it transfers right after a prompt-only resize, so valid_only
    skips the page-less tail). Used to exercise the SWA + MTP transfer path.
    """
    kv_cache = mgr._create_kv_cache(request_id, None, None)
    if not mgr.enable_block_reuse:
        kv_cache.stop_committing()
    success = kv_cache.resume(torch.cuda.current_stream().cuda_stream)
    assert success, f"Failed to resume kv_cache for request {request_id}"
    if is_generation:
        kv_cache.enable_swa_scratch_reuse = False
        kv_cache.resize(prompt_len + scratch_tokens, history_length=prompt_len)
    else:
        kv_cache.resize(prompt_len)
        kv_cache.resize(None, history_length=prompt_len)
        kv_cache.enable_swa_scratch_reuse = False
    return kv_cache


# ---------------------------------------------------------------------------
# Transceiver creation (threaded)
# ---------------------------------------------------------------------------
def _create_transceiver_in_thread(
    rank, mapping, cache_manager, dist_mock, attention_type, config, results, errors
):
    """Thread target: create one KvCacheTransceiverV2."""
    try:
        tc = KvCacheTransceiverV2(
            mapping=mapping,
            dist=dist_mock,
            kv_cache_manager=cache_manager,
            cache_transceiver_config=config,
        )
        results[rank] = tc
    except Exception as e:
        errors[rank] = e


def create_instance_transceivers(
    tp: int,
    pp: int,
    enable_dp: bool,
    cache_managers: List,
    config: CacheTransceiverConfig,
    is_mla: bool,
    cp_size: int = 1,
) -> List[KvCacheTransceiverV2]:
    """Create KvCacheTransceiverV2 for all ranks via threaded init.

    cp_size > 1 builds a helix-CP instance (see _create_managers_for_instance).
    Ranks are enumerated (pp_rank * tp + tp_rank) * cp_size + cp_rank, matching
    the manager order so cache_managers[rank] lines up.
    HELIX VALIDATE: the ThreadSafeDistributed collective grouping below is the
    tp*pp structure; cp ranks are extra world ranks. Confirm the transceiver's
    setup collectives behave under helix (they may need cp-aware grouping).
    """
    world_size = tp * pp * cp_size
    shared = {
        "barrier": threading.Barrier(world_size),
        # Per-group barriers for the grouped collectives, so attention-DP instances can
        # diverge in how many collectives they issue: pp_allgather syncs the PP ranks that
        # share a tp_rank, tp_allgather the TP ranks that share a pp_rank.
        "pp_barriers": [threading.Barrier(pp) for _ in range(tp * cp_size)],
        "tp_barriers": [threading.Barrier(tp) for _ in range(pp * cp_size)],
        "lock": threading.Lock(),
    }
    results = [None] * world_size
    errors = [None] * world_size
    threads = []
    cp_config = {"cp_type": "HELIX"} if cp_size > 1 else None

    attention_type = AttentionTypeCpp.MLA if is_mla else AttentionTypeCpp.DEFAULT

    for rank in range(world_size):
        cp_rank = rank % cp_size
        base = rank // cp_size
        pp_rank = base // tp
        tp_rank = base % tp
        mapping = Mapping(
            world_size=world_size,
            rank=rank,
            tp_size=tp,
            pp_size=pp,
            cp_size=cp_size,
            cp_config=cp_config,
            enable_attention_dp=enable_dp,
        )
        dist_mock = ThreadSafeDistributed(
            rank, world_size, tp, pp, tp_rank, pp_rank, shared, cp_rank, cp_size
        )
        t = threading.Thread(
            target=_create_transceiver_in_thread,
            args=(
                rank,
                mapping,
                cache_managers[rank],
                dist_mock,
                attention_type,
                config,
                results,
                errors,
            ),
        )
        threads.append(t)

    for t in threads:
        t.start()
    for t in threads:
        t.join()

    for rank, err in enumerate(errors):
        if err is not None:
            raise err

    return results


# ---------------------------------------------------------------------------
# Verification helpers
# ---------------------------------------------------------------------------
def _get_block_data_for_layer(
    mgr, request_id, layer_idx, use_v2, expected_valid=None, drop_tail_blocks=0
):
    """Get block data for a specific layer and request in HND layout.

    Returns tensor of shape [num_blocks, kv_factor, num_kv_heads_per_rank, tokens_per_block, head_dim].
    The pool's physical layout is HND (heads vary slowest within each kv_factor),
    so we must use kv_layout='HND' to get a correct view.  Using the default NHD
    layout would reinterpret HND memory as NHD, which breaks cross-TP verification
    because the formatter concat/split operates on HND-ordered flat data.

    *drop_tail_blocks* removes that many blocks from the tail before anything
    else -- the MTP speculative/scratch blocks, which sit past the prompt and are
    never transferred. It must be a COUNT from the tail (not a prompt_blocks
    cut): the SWA front may already be evicted, so ``len(valid_indices)`` is not
    prompt_blocks and a positional cut would leave the scratch in and drop a live
    window block. Gen side only (the ctx cache has no scratch).

    If *expected_valid* is given, keeps only the last *expected_valid* blocks
    (used for sliding-window: only the window-tail blocks were transferred).
    This works regardless of whether eviction has happened:
      - Pre-eviction: all blocks present, keep last N.
      - Post-eviction: stale blocks already removed (or -1), keep last N is no-op.
    """
    block_indices = mgr.get_batch_cache_indices([request_id], layer_idx)[0]
    valid_indices = [idx for idx in block_indices if idx >= 0]
    if not valid_indices:
        return None
    if drop_tail_blocks > 0 and len(valid_indices) > drop_tail_blocks:
        valid_indices = valid_indices[:-drop_tail_blocks]
    if expected_valid is not None and len(valid_indices) > expected_valid:
        valid_indices = valid_indices[-expected_valid:]
    layer_buffer = mgr.get_buffers(layer_idx, kv_layout="HND")
    return layer_buffer[valid_indices]


def _gather_full_layer_data(
    managers: List,
    tp: int,
    pp: int,
    request_id: int,
    layer_idx: int,
    use_v2: bool,
    num_kv_heads: int,
    expected_valid: Optional[int] = None,
    enable_dp: bool = False,
    req_idx: int = 0,
    num_layers: int = NUM_LAYERS,
    drop_tail_blocks: int = 0,
):
    """Gather the full (unsharded) KV data for a layer by concatenating across TP ranks.

    Returns tensor of shape [num_blocks, kv_factor, num_kv_heads, tokens_per_block, head_dim]
    (HND layout) where num_kv_heads is the total (not per-rank) count.
    For replicated heads (MLA), returns data from a single TP rank (all are identical).
    """
    pp_rank = _pp_rank_of_layer(layer_idx, num_layers, pp)
    is_replicated = num_kv_heads < tp

    if enable_dp:
        # Only one TP rank handled this request
        tp_rank = req_idx % tp
        rank = pp_rank * tp + tp_rank
        return _get_block_data_for_layer(
            managers[rank], request_id, layer_idx, use_v2, expected_valid, drop_tail_blocks
        )

    if is_replicated:
        # All TP ranks have identical data; take from tp_rank=0
        rank = pp_rank * tp + 0
        return _get_block_data_for_layer(
            managers[rank], request_id, layer_idx, use_v2, expected_valid, drop_tail_blocks
        )

    # Gather from all TP ranks and concat along kv_heads dim.
    # In HND layout the shape is [blocks, kv_factor, heads, tpb, dim],
    # so heads are at dim=2.
    tp_data = []
    for tp_rank in range(tp):
        rank = pp_rank * tp + tp_rank
        data = _get_block_data_for_layer(
            managers[rank], request_id, layer_idx, use_v2, expected_valid, drop_tail_blocks
        )
        if data is not None:
            tp_data.append(data)

    if not tp_data:
        return None
    return torch.cat(tp_data, dim=2)


def _get_layer_to_window_size(
    managers: List,
    tp: int,
    pp: int,
    use_v2: bool,
    num_layers: int = NUM_LAYERS,
) -> Dict[int, Optional[int]]:
    """Build a mapping from global layer_id to its sliding_window_size.

    For V2 managers, extracts from page_table.layer_group_metas.

    Aggregates from all PP ranks since each rank only handles a subset of layers.
    """
    layer_to_window: Dict[int, Optional[int]] = {}

    # Get one manager per PP rank (tp_rank=0 for each pp_rank)
    for pp_rank in range(pp):
        manager = managers[pp_rank * tp]

        page_table = KVRegionExtractorV1(manager).page_table
        for lg in page_table.layer_groups:
            for layer_id in get_global_layer_ids(lg):
                layer_to_window[layer_id] = lg.sliding_window_size

    return layer_to_window


def verify_all_requests(
    request_lengths: List[int],
    ctx_managers: List,
    gen_managers: List,
    ctx_tp: int,
    ctx_pp: int,
    gen_tp: int,
    gen_pp: int,
    ctx_enable_dp: bool,
    gen_enable_dp: bool,
    ctx_request_ids: List[int],
    gen_request_ids: List[int],
    is_mla: bool,
    use_v2: bool,
    max_attention_window_vec: Optional[List[int]] = None,
    num_layers: int = NUM_LAYERS,
    gen_scratch_tokens: int = 0,
):
    """Verify transferred cache data for all requests.

    For each request and each global layer, reconstruct the full (unsharded)
    KV data from both ctx and gen sides by gathering across TP ranks, then
    compare.  This correctly handles asymmetric TP/PP configurations.
    """
    num_kv_heads = 1 if is_mla else NUM_KV_HEADS

    # Build layer_to_window_size mapping from ctx_managers
    # Aggregate from all PP ranks since each rank only handles a subset of layers
    layer_to_window = _get_layer_to_window_size(ctx_managers, ctx_tp, ctx_pp, use_v2, num_layers)

    for req_idx, req_len in enumerate(request_lengths):
        ctx_rid = ctx_request_ids[req_idx]
        gen_rid = gen_request_ids[req_idx]

        # MTP speculative/scratch blocks on the gen side: sized past the prompt,
        # never transferred. Drop them (from the tail) before comparing. The ctx
        # cache has no scratch. Count = ceil((prompt+scratch)/tpb) - ceil(prompt/tpb).
        prompt_blocks = (req_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK
        gen_scratch_blocks = (
            req_len + gen_scratch_tokens + TOKENS_PER_BLOCK - 1
        ) // TOKENS_PER_BLOCK - prompt_blocks

        for layer_idx in range(num_layers):
            # Compute expected_valid: the number of non-stale blocks that
            # were actually transferred, using the same eviction formula as
            # _create_chunk.  Only compare these blocks in verification.
            expected_valid = None
            win = layer_to_window.get(layer_idx)
            if win is not None and win < MAX_SEQ_LEN:
                stale_end = max(0, (req_len + 1 - win) // TOKENS_PER_BLOCK)
                expected_valid = prompt_blocks - stale_end

            ctx_full = _gather_full_layer_data(
                ctx_managers,
                ctx_tp,
                ctx_pp,
                ctx_rid,
                layer_idx,
                use_v2,
                num_kv_heads,
                expected_valid,
                ctx_enable_dp,
                req_idx,
                num_layers,
                0,  # ctx has no MTP scratch tail
            )
            gen_full = _gather_full_layer_data(
                gen_managers,
                gen_tp,
                gen_pp,
                gen_rid,
                layer_idx,
                use_v2,
                num_kv_heads,
                expected_valid,
                gen_enable_dp,
                req_idx,
                num_layers,
                gen_scratch_blocks,
            )

            if ctx_full is None or gen_full is None:
                continue

            assert ctx_full.shape == gen_full.shape, (
                f"Shape mismatch at req={req_idx} layer={layer_idx}: "
                f"ctx={ctx_full.shape} gen={gen_full.shape}"
            )
            torch.testing.assert_close(
                gen_full,
                ctx_full,
                rtol=0,
                atol=0,
                msg=lambda m: f"Data mismatch at req={req_idx} layer={layer_idx}: {m}",
            )


def verify_all_requests_helix(
    request_lengths: List[int],
    ctx_managers: List,
    gen_managers: List,
    ctx_request_ids: List[int],
    gen_request_ids: List[int],
    is_mla: bool,
    gen_cp: int,
    max_attention_window_vec: Optional[List[int]] = None,
    num_layers: int = NUM_LAYERS,
):
    """Verify a helix-CP transfer against the ctx block stride.

    ctx holds the whole prompt (cp_size=1); each generation CP rank owns the
    global blocks strided by [cp_rank::cp_size].

    Restricted to ctx_tp==ctx_pp==gen_tp==gen_pp==1 so a rank is exactly a CP
    rank (gen rank r == cp_rank r, ctx rank 0). That keeps the comparison a plain
    positional stride over the ctx block list, matching transfer.py:1036
    (``src_block_ids[cp_rank::cp_size]``) without also folding in TP head-concat.

    HELIX VALIDATE: the ledger stride below assumes the sender strides over the
    *transferred* (windowed) ctx block list in positional order. Confirm on GPU
    that the SWA window + stride interaction matches (no_window is the clean case).
    """
    num_kv_heads = 1 if is_mla else NUM_KV_HEADS
    # ctx side is cp_size=1: rank 0 holds the whole global block list.
    layer_to_window = _get_layer_to_window_size(ctx_managers, 1, 1, True, num_layers)

    for req_idx, req_len in enumerate(request_lengths):
        ctx_rid = ctx_request_ids[req_idx]
        gen_rid = gen_request_ids[req_idx]
        prompt_blocks = (req_len + TOKENS_PER_BLOCK - 1) // TOKENS_PER_BLOCK

        for layer_idx in range(num_layers):
            expected_valid = None
            win = layer_to_window.get(layer_idx)
            if win is not None and win < MAX_SEQ_LEN:
                stale_end = max(0, (req_len + 1 - win) // TOKENS_PER_BLOCK)
                expected_valid = prompt_blocks - stale_end

            # ctx: the full (global) transferred block set, positional order.
            ctx_full = _get_block_data_for_layer(
                ctx_managers[0], ctx_rid, layer_idx, True, expected_valid
            )
            if ctx_full is None:
                continue

            for cp_rank in range(gen_cp):
                gen_full = _get_block_data_for_layer(
                    gen_managers[cp_rank], gen_rid, layer_idx, True
                )
                # Ledger stride: rank r owns global blocks [r::cp_size].
                expected = ctx_full[cp_rank::gen_cp]
                # The gen ledger rounds capacity up to whole ledger blocks, so a
                # rank that owns no token in the trailing ledger block still binds
                # one extra (unfilled) page past its owned blocks. The transceiver
                # only fills the rank-local prompt (the leading `n` pages), so
                # compare just those -- exactly the ctx stride count.
                n = expected.shape[0]
                if gen_full is not None and gen_full.shape[0] > n:
                    assert gen_full.shape[0] <= n + 1, (
                        f"cp_rank {cp_rank} holds {gen_full.shape[0]} blocks but "
                        f"the ctx stride is {n} (at most one trailing ledger page "
                        f"is expected) at req={req_idx} layer={layer_idx}"
                    )
                    gen_full = gen_full[:n]
                if n == 0:
                    assert gen_full is None or gen_full.shape[0] == 0, (
                        f"cp_rank {cp_rank} holds blocks but ctx stride is empty "
                        f"at req={req_idx} layer={layer_idx}"
                    )
                    continue
                assert gen_full is not None, (
                    f"cp_rank {cp_rank} missing blocks at req={req_idx} layer={layer_idx}"
                )
                assert gen_full.shape == expected.shape, (
                    f"Helix shape mismatch at req={req_idx} layer={layer_idx} "
                    f"cp_rank={cp_rank}: gen={gen_full.shape} "
                    f"ctx[{cp_rank}::{gen_cp}]={expected.shape}"
                )
                torch.testing.assert_close(
                    gen_full,
                    expected,
                    rtol=0,
                    atol=0,
                    msg=lambda m: (
                        f"Helix data mismatch at req={req_idx} layer={layer_idx} "
                        f"cp_rank={cp_rank}: {m}"
                    ),
                )
    _ = num_kv_heads  # (kept for signature parity with verify_all_requests)


# ---------------------------------------------------------------------------
# Main test orchestrator
# ---------------------------------------------------------------------------
def _send_prefill_in_pieces(transceiver, req: LlmRequest, piece_tokens: int) -> int:
    """Send the prompt's KV in pieces of piece_tokens and return the number of pieces sent.

    One context chunk per piece, as the executor does with enable_pipelined_transfer.
    """
    # A non-final piece without a complete block is not sent, which would break the count.
    assert piece_tokens % TOKENS_PER_BLOCK == 0
    pieces = 0
    for start in range(0, req.prompt_len, piece_tokens):
        end = min(start + piece_tokens, req.prompt_len)
        req.py_last_context_chunk = (start, end)
        req.context_current_position = end
        transceiver.respond_and_send_async(req)
        pieces += 1
    return pieces


def run_transfer_test(
    ctx_tp: int,
    ctx_pp: int,
    gen_tp: int,
    gen_pp: int,
    ctx_enable_dp: bool,
    gen_enable_dp: bool,
    is_mla: bool,
    use_v2: bool,
    max_attention_window_vec: Optional[List[int]] = None,
    num_layers: int = NUM_LAYERS,
    request_lengths: Optional[List[int]] = None,
    expect_cpp_bounce: bool = False,
    gen_scratch_tokens: int = 0,
    gen_cp: int = 1,
    pipelined_piece_tokens: Optional[int] = None,
    num_draft_layers: int = 0,
) -> None:
    """Run a full KV transfer test using KvCacheTransceiverV2.

    num_draft_layers > 0 adds K-only draft-owned layers after the target's on
    both sides (a unified draft KV cache); they must transfer like the target's.

    pipelined_piece_tokens enables enable_pipelined_transfer: the context sends each prompt in
    pieces of that many tokens instead of whole.

    gen_scratch_tokens > 0 sizes the generation cache to
    prompt_len + gen_scratch_tokens (history still prompt_len), materialising an
    MTP speculative tail block past the prompt. Combined with a sliding window
    this exercises the SWA + MTP transfer path -- the combination that regressed
    GSM8K (bug 6676406) and that no prior case covered.

    gen_cp > 1 builds a helix-CP generation instance (context stays cp_size=1,
    i.e. it holds the whole global prompt). Each generation CP rank owns the
    global blocks strided by [cp_rank::cp_size] -- the ledger model where a
    global block b lives on rank ``b %% cp_size`` (see transfer.py:1036 and
    KVCacheManagerV2._helix_local_len). Helix rejects draft tokens, so this is
    incompatible with gen_scratch_tokens > 0.
    HELIX VALIDATE: the whole gen_cp path (dist-mock collectives, ledger block
    ownership, strided verify) is exercised only on GPU -- validate end to end.
    """
    if request_lengths is None:
        request_lengths = REQUEST_LENGTHS
    assert not (gen_cp > 1 and gen_scratch_tokens > 0), (
        "helix rejects draft/MTP tokens; gen_cp>1 requires gen_scratch_tokens==0"
    )
    max_batch_size = max(MAX_BATCH_SIZE, len(request_lengths))
    ctx_world = ctx_tp * ctx_pp
    gen_world = gen_tp * gen_pp * gen_cp

    # 1. Create cache managers
    ctx_managers = _create_managers_for_instance(
        ctx_tp,
        ctx_pp,
        ctx_enable_dp,
        is_mla,
        use_v2,
        max_attention_window_vec,
        num_layers,
        max_batch_size,
        num_draft_layers=num_draft_layers,
    )
    gen_managers = _create_managers_for_instance(
        gen_tp,
        gen_pp,
        gen_enable_dp,
        is_mla,
        use_v2,
        max_attention_window_vec,
        num_layers,
        max_batch_size,
        cp_size=gen_cp,
        num_draft_layers=num_draft_layers,
    )

    # 2. Initialize data: random for ctx, zeros for gen
    _init_pool_data(ctx_managers, ctx_tp, is_mla, use_v2, fill_random=True, seed_base=1000)
    # Gen pools start zeroed; gen_cp folds into the TP stride so per-rank pool
    # seeding stays a no-op (fill_random=False ignores the seed anyway).
    _init_pool_data(gen_managers, gen_tp * gen_cp, is_mla, use_v2, fill_random=False)

    # 3. Create KvCacheTransceiverV2 instances (threaded init)
    config = CacheTransceiverConfig(
        backend="NIXL",
        transceiver_runtime="PYTHON",
        max_tokens_in_buffer=512,
        # Bounce is off unless expect_cpp_bounce. Then a positive kv_cache_bounce_size_mb
        # enables the C++ agent bounce, and agent_bounce_params (a loose gate) exercises
        # the config plumbing end to end.
        kv_cache_bounce_size_mb=64 if expect_cpp_bounce else 0,
        agent_bounce_params={
            "min_descriptor_count": "1",
            "max_average_descriptor_size": "1MB",
        }
        if expect_cpp_bounce
        else None,
        enable_pipelined_transfer=pipelined_piece_tokens is not None,
    )
    ctx_tcs = create_instance_transceivers(
        ctx_tp, ctx_pp, ctx_enable_dp, ctx_managers, config, is_mla
    )
    gen_tcs = create_instance_transceivers(
        gen_tp, gen_pp, gen_enable_dp, gen_managers, config, is_mla, cp_size=gen_cp
    )

    # Defined before the try so the finally-block teardown can always close them,
    # even if setup raises before any sequence is added.
    ctx_kv_caches: Dict[int, List] = {r: [] for r in range(ctx_world)}
    gen_kv_caches: Dict[int, List] = {r: [] for r in range(gen_world)}

    try:
        ctx_info_endpoint = ctx_tcs[0]._context_info_endpoint

        # ===== 4. Create requests =====
        ctx_handle_map: Dict[int, List] = {r: [] for r in range(ctx_world)}
        gen_handle_map: Dict[int, List] = {r: [] for r in range(gen_world)}
        ctx_request_ids: List[int] = []
        gen_request_ids: List[int] = []

        sampling_params = SamplingParams()

        for req_idx, req_len in enumerate(request_lengths):
            unique_rid = uuid.uuid4().int & 0x7FFFFFFFFFFFFFFF
            ctx_rid = req_idx * 2
            gen_rid = req_idx * 2 + 1
            ctx_request_ids.append(ctx_rid)
            gen_request_ids.append(gen_rid)

            ctx_dp_rank = req_idx % ctx_tp if ctx_enable_dp else 0

            ctx_request = LlmRequest(
                request_id=ctx_rid,
                max_new_tokens=1,
                input_tokens=list(range(req_len)),
                sampling_config=tensorrt_llm.bindings.SamplingConfig(
                    sampling_params._get_sampling_config()
                ),
                is_streaming=False,
                llm_request_type=LlmRequestType.LLMREQUEST_TYPE_CONTEXT_ONLY,
            )
            ctx_request.py_disaggregated_params = DisaggregatedParams(disagg_request_id=unique_rid)

            gen_request = LlmRequest(
                request_id=gen_rid,
                max_new_tokens=1,
                input_tokens=list(range(req_len)),
                sampling_config=tensorrt_llm.bindings.SamplingConfig(
                    sampling_params._get_sampling_config()
                ),
                is_streaming=False,
                llm_request_type=LlmRequestType.LLMREQUEST_TYPE_GENERATION_ONLY,
            )
            gen_request.py_disaggregated_params = DisaggregatedParams(
                ctx_request_id=ctx_rid,
                ctx_dp_rank=ctx_dp_rank,
                ctx_info_endpoint=ctx_info_endpoint,
                disagg_request_id=unique_rid,
            )
            if gen_scratch_tokens > 0:
                # Mark the request as MTP so its draft length matches the
                # materialised speculative tail (see _add_sequence).
                gen_request.py_draft_tokens = list(range(gen_scratch_tokens))

            for rank in range(ctx_world):
                tp_rank = rank % ctx_tp
                should_handle = (not ctx_enable_dp) or (req_idx % ctx_tp == tp_rank)
                if should_handle:
                    ctx_handle_map[rank].append((req_idx, ctx_request))

            for rank in range(gen_world):
                # rank = (pp_rank * gen_tp + tp_rank) * gen_cp + cp_rank.
                # Every CP rank of a handling TP rank receives its strided slice.
                tp_rank = (rank // gen_cp) % gen_tp
                should_handle = (not gen_enable_dp) or (req_idx % gen_tp == tp_rank)
                if not should_handle:
                    continue
                if gen_cp > 1:
                    # Helix: each CP rank gets its own request. The transceiver
                    # bounds the transferred block list by prompt_len, which must
                    # be the *rank-local* strided length so it matches the ctx
                    # sender's src_block_ids[cp_rank::cp_size]. total_input_len_cp
                    # carries the global length the manager sizes the ledger off.
                    # A shared request cannot hold a per-rank prompt_len.
                    cp_rank = rank % gen_cp
                    req = _make_gen_request(
                        gen_rid,
                        req_len,
                        unique_rid,
                        ctx_rid,
                        ctx_dp_rank,
                        ctx_info_endpoint,
                        sampling_params,
                    )
                    local_len = _helix_local_len(req_len, cp_rank, gen_cp)
                    req.prompt_len = local_len
                    req.py_prompt_len = local_len
                    req.total_input_len_cp = req_len
                    gen_handle_map[rank].append((req_idx, req))
                else:
                    gen_handle_map[rank].append((req_idx, gen_request))

        # 5. Add sequences and gen receive first
        for rank in range(gen_world):
            for req_idx, req in gen_handle_map[rank]:
                # Under helix req.prompt_len is the rank-local strided length
                # (for the transceiver); the manager ledger must be sized off the
                # GLOBAL length, so resize off request_lengths[req_idx] instead.
                resize_len = request_lengths[req_idx] if gen_cp > 1 else req.prompt_len
                kv = _add_sequence(
                    gen_managers[rank],
                    req.py_request_id,
                    resize_len,
                    use_v2,
                    is_generation=True,
                    scratch_tokens=gen_scratch_tokens,
                )
                if kv is not None:
                    gen_kv_caches[rank].append(kv)
                gen_tcs[rank].request_and_receive_async(req)

        # 6. ctx send after
        pieces_sent = 0
        for rank in range(ctx_world):
            for req_idx, req in ctx_handle_map[rank]:
                kv = _add_sequence(ctx_managers[rank], req.py_request_id, req.prompt_len, use_v2)
                if kv is not None:
                    ctx_kv_caches[rank].append(kv)
                if pipelined_piece_tokens is None:
                    ctx_tcs[rank].respond_and_send_async(req)
                else:
                    pieces_sent += _send_prefill_in_pieces(
                        ctx_tcs[rank], req, pipelined_piece_tokens
                    )

        # 7. Wait for completion (threaded, dist calls inside)
        run_concurrent(
            ctx_tcs, lambda tc: tc.check_context_transfer_status(None, mark_complete=True)
        )
        run_concurrent(gen_tcs, lambda tc: tc.check_gen_transfer_status(None))

        # 8. Verify
        if gen_cp > 1:
            verify_all_requests_helix(
                request_lengths=request_lengths,
                ctx_managers=ctx_managers,
                gen_managers=gen_managers,
                ctx_request_ids=ctx_request_ids,
                gen_request_ids=gen_request_ids,
                is_mla=is_mla,
                gen_cp=gen_cp,
                max_attention_window_vec=max_attention_window_vec,
                num_layers=num_layers + num_draft_layers,
            )
        else:
            verify_all_requests(
                request_lengths=request_lengths,
                ctx_managers=ctx_managers,
                gen_managers=gen_managers,
                ctx_tp=ctx_tp,
                ctx_pp=ctx_pp,
                gen_tp=gen_tp,
                gen_pp=gen_pp,
                ctx_enable_dp=ctx_enable_dp,
                gen_enable_dp=gen_enable_dp,
                ctx_request_ids=ctx_request_ids,
                gen_request_ids=gen_request_ids,
                is_mla=is_mla,
                use_v2=use_v2,
                max_attention_window_vec=max_attention_window_vec,
                num_layers=num_layers,
                gen_scratch_tokens=gen_scratch_tokens,
            )

        # Programmatic bounce check (no log parsing): every agent must have the C++
        # bounce v2 transport active, and at least one transfer must have been routed
        # through it (only KV senders submit WRITEs, so we assert on the total).
        if expect_cpp_bounce:
            agents = [tc._transfer_worker._agent for tc in ctx_tcs + gen_tcs]
            assert all(getattr(agent, "bounce_enabled", False) for agent in agents), (
                "C++ bounce v2 transport is not active on every NIXL agent"
            )
            total_bounce_submits = sum(agent.bounce_submit_count for agent in agents)
            assert total_bounce_submits > 0, "no transfer was routed through the bounce fast path"
            if pipelined_piece_tokens is not None:
                # Exact for one context rank writing one generation rank: each piece is one
                # transfer request, and every one of them must take the bounce.
                assert len(ctx_tcs) == len(gen_tcs) == 1
                ctx_agents = [tc._transfer_worker._agent for tc in ctx_tcs]
                assert all(agent.bounce_reject_count == 0 for agent in ctx_agents), [
                    agent.bounce_reject_counts for agent in ctx_agents
                ]
                assert sum(agent.bounce_submit_count for agent in ctx_agents) == pieces_sent

    finally:
        # Teardown must run on the failure path too, in this order:
        #   drain GPU work -> close kv caches -> shutdown transceivers.
        # If verify raises, skipping the drain/close and tearing down the NIXL
        # agents while registered memory / in-flight transfers are still live
        # makes the C++ side abort during GC and leaves the messenger/task-queue
        # listener threads hanging (process never exits). Each step is guarded so
        # a teardown error cannot mask the real test failure.
        try:
            torch.cuda.current_stream().synchronize()
        except Exception:
            pass
        for rank in range(ctx_world):
            for kv in ctx_kv_caches[rank]:
                try:
                    kv.close()
                except Exception:
                    pass
        for rank in range(gen_world):
            for kv in gen_kv_caches[rank]:
                try:
                    kv.close()
                except Exception:
                    pass
        for tc in ctx_tcs + gen_tcs:
            try:
                tc.shutdown()
            except Exception:
                pass
        # On the failure path pytest keeps the traceback alive, which pins this
        # frame's locals (ctx_tcs/gen_tcs) and defers finalizing their NIXL C++
        # agents to a later GC -- which then lands inside the *next* test's NIXL
        # init and aborts the process. Drop the references and collect now so the
        # C++ teardown happens deterministically here, within this test.
        tc = None
        ctx_tcs = gen_tcs = None
        try:
            gc.collect()
            torch.cuda.synchronize()
        except Exception:
            pass


# ---------------------------------------------------------------------------
# Test configurations
# ---------------------------------------------------------------------------
TEST_CONFIGS = [
    # (ctx_tp, ctx_pp, gen_tp, gen_pp, ctx_dp, gen_dp, is_mla, test_id)
    # --- Basic symmetric ---
    (1, 1, 1, 1, False, False, False, "tp1_pp1"),
    (2, 1, 2, 1, False, False, False, "tp2_pp1"),
    (1, 2, 1, 2, False, False, False, "tp1_pp2"),
    (2, 2, 2, 2, False, False, False, "tp2_pp2"),
    (4, 1, 4, 1, False, False, False, "tp4_pp1"),
    (1, 4, 1, 4, False, False, False, "tp1_pp4"),
    (4, 2, 4, 2, False, False, False, "tp4_pp2"),
    (2, 4, 2, 4, False, False, False, "tp2_pp4"),
    (4, 4, 4, 4, False, False, False, "tp4_pp4"),
    # --- Asymmetric TP/PP between ctx and gen ---
    (1, 2, 1, 1, False, False, False, "pp2_to_pp1"),
    (2, 1, 1, 2, False, False, False, "tp2_to_pp2"),
    (4, 1, 2, 2, False, False, False, "tp4_to_tp2pp2"),
    (4, 1, 1, 4, False, False, False, "tp4_to_pp4"),
    (2, 2, 4, 1, False, False, False, "tp2pp2_to_tp4"),
    (1, 4, 4, 1, False, False, False, "pp4_to_tp4"),
    (4, 2, 2, 4, False, False, False, "tp4pp2_to_tp2pp4"),
    # --- DP (attention data parallelism) ---
    (2, 1, 2, 1, True, True, False, "tp2_dp_both"),
    (4, 1, 4, 1, True, True, False, "tp4_dp_both"),
    (2, 1, 1, 2, True, False, False, "ctx_dp_gen_pp2"),
    (4, 1, 2, 2, True, False, False, "tp4_dp_to_tp2pp2"),
    (2, 2, 2, 2, True, True, False, "tp2_pp2_dp_both"),
    (4, 2, 4, 2, True, True, False, "tp4_pp2_dp_both"),
    # --- MLA (Multi-head Latent Attention / SELFKONLY) ---
    (1, 1, 1, 1, False, False, True, "mla_tp1_pp1"),
    (2, 1, 2, 1, False, False, True, "mla_tp2"),
    (4, 1, 4, 1, False, False, True, "mla_tp4"),
    (2, 2, 2, 2, False, False, True, "mla_tp2_pp2"),
    (4, 2, 4, 2, False, False, True, "mla_tp4_pp2"),
    (2, 1, 2, 1, True, False, True, "mla_tp2_ctx_dp"),
    (4, 1, 4, 1, True, True, True, "mla_tp4_dp_both"),
    (2, 1, 1, 2, False, False, True, "mla_tp2_to_pp2"),
    # --- MLA asymmetric TP ---
    (4, 1, 8, 1, True, False, True, "mla_ctx_dp4_gen_tp8"),
    (4, 1, 8, 1, False, False, True, "mla_ctx_tp4_gen_tp8"),
]

WINDOW_CONFIGS = [
    # (max_attention_window_vec, test_id)
    (None, "no_window"),
    ([24], "uniform_window"),
    ([MAX_SEQ_LEN, 24], "vswa"),
]

# --- Uneven PP layer configs ---
# These test configurations use num_layers not evenly divisible by pp_size,
# so the first (num_layers % pp_size) PP ranks get one extra layer.
# Mirrors the C++ getLayerNumPPRank pattern in cacheTransceiverTest.cpp.
# Format: (ctx_tp, ctx_pp, gen_tp, gen_pp, ctx_dp, gen_dp, is_mla, num_layers, test_id)
UNEVEN_PP_CONFIGS = [
    # 5 layers, pp=2 → [3, 2]
    (1, 2, 1, 2, False, False, False, 5, "5L_tp1_pp2"),
    (2, 2, 2, 2, False, False, False, 5, "5L_tp2_pp2"),
    # 5 layers, pp=3 → [2, 2, 1]
    (1, 3, 1, 3, False, False, False, 5, "5L_tp1_pp3"),
    # 7 layers, pp=2 → [4, 3]
    (1, 2, 1, 2, False, False, False, 7, "7L_tp1_pp2"),
    (2, 2, 2, 2, False, False, False, 7, "7L_tp2_pp2"),
    # 7 layers, pp=3 → [3, 2, 2]
    (1, 3, 1, 3, False, False, False, 7, "7L_tp1_pp3"),
    # 7 layers, pp=4 → [2, 2, 2, 1]
    (1, 4, 1, 4, False, False, False, 7, "7L_tp1_pp4"),
    # Asymmetric TP/PP with uneven layers
    (2, 1, 1, 2, False, False, False, 5, "5L_tp2_to_pp2"),
    (1, 2, 2, 1, False, False, False, 5, "5L_pp2_to_tp2"),
    (4, 1, 1, 4, False, False, False, 7, "7L_tp4_to_pp4"),
    (1, 4, 4, 1, False, False, False, 7, "7L_pp4_to_tp4"),
    (2, 2, 1, 4, False, False, False, 7, "7L_tp2pp2_to_pp4"),
    # Uneven layers + MLA
    (1, 2, 1, 2, False, False, True, 5, "5L_mla_tp1_pp2"),
    (2, 2, 2, 2, False, False, True, 5, "5L_mla_tp2_pp2"),
    (1, 3, 1, 3, False, False, True, 7, "7L_mla_tp1_pp3"),
]


# ---------------------------------------------------------------------------
# pytest test functions
# ---------------------------------------------------------------------------
@pytest.mark.timeout(180)
@pytest.mark.parametrize(
    "ctx_tp,ctx_pp,gen_tp,gen_pp,ctx_enable_dp,gen_enable_dp,is_mla",
    [(c[0], c[1], c[2], c[3], c[4], c[5], c[6]) for c in TEST_CONFIGS],
    ids=[c[7] for c in TEST_CONFIGS],
)
@pytest.mark.parametrize(
    "max_attention_window_vec",
    [w[0] for w in WINDOW_CONFIGS],
    ids=[w[1] for w in WINDOW_CONFIGS],
)
@pytest.mark.parametrize("use_v2", [True], ids=["v2"])
@pytest.mark.parametrize("gen_scratch_tokens", [0, TOKENS_PER_BLOCK], ids=["spec_off", "spec_on"])
def test_cache_transceiver(
    ctx_tp,
    ctx_pp,
    gen_tp,
    gen_pp,
    ctx_enable_dp,
    gen_enable_dp,
    is_mla,
    max_attention_window_vec,
    use_v2,
    gen_scratch_tokens,
):
    """Test KvCacheTransceiverV2 with V2 cache managers.

    The ``gen_scratch_tokens`` dimension (id ``spec_off`` / ``spec_on``) adds a
    speculative-decoding tail: the generation cache is sized
    prompt_len + TOKENS_PER_BLOCK so a surplus block sits past the prompt,
    exercising the manager+transceiver combinations *with speculation* -- the
    tail is generic (any spec mode, e.g. MTP or Eagle), including the
    sliding-window case that regressed GSM8K under MTP (bug 6676406). Every prior
    run was spec_off (gen_scratch_tokens=0). The transfer must still cover exactly
    the prompt blocks the receiver needs; a count-based slice would let the
    surplus shift or drop a live prompt block.
    """
    # Speculative scratch materialisation (resize past history_length) is V2-only

    print(
        f"\nRunning transfer test: "
        f"ctx_tp={ctx_tp} ctx_pp={ctx_pp} gen_tp={gen_tp} gen_pp={gen_pp} "
        f"ctx_dp={ctx_enable_dp} gen_dp={gen_enable_dp} "
        f"mla={is_mla} v2={use_v2} window={max_attention_window_vec} "
        f"scratch={gen_scratch_tokens}"
    )

    run_transfer_test(
        ctx_tp=ctx_tp,
        ctx_pp=ctx_pp,
        gen_tp=gen_tp,
        gen_pp=gen_pp,
        ctx_enable_dp=ctx_enable_dp,
        gen_enable_dp=gen_enable_dp,
        is_mla=is_mla,
        use_v2=use_v2,
        max_attention_window_vec=max_attention_window_vec,
        gen_scratch_tokens=gen_scratch_tokens,
    )

    print("PASSED")


@pytest.mark.timeout(180)
@pytest.mark.parametrize(
    "ctx_tp,ctx_pp,gen_tp,gen_pp,ctx_enable_dp,gen_enable_dp,is_mla,num_layers",
    [(c[0], c[1], c[2], c[3], c[4], c[5], c[6], c[7]) for c in UNEVEN_PP_CONFIGS],
    ids=[c[8] for c in UNEVEN_PP_CONFIGS],
)
@pytest.mark.parametrize(
    "max_attention_window_vec",
    [w[0] for w in WINDOW_CONFIGS],
    ids=[w[1] for w in WINDOW_CONFIGS],
)
@pytest.mark.parametrize("use_v2", [True], ids=["v2"])
def test_cache_transceiver_uneven_pp(
    ctx_tp,
    ctx_pp,
    gen_tp,
    gen_pp,
    ctx_enable_dp,
    gen_enable_dp,
    is_mla,
    num_layers,
    max_attention_window_vec,
    use_v2,
):
    """Test KvCacheTransceiverV2 with uneven layers-per-PP-rank.

    When num_layers is not evenly divisible by pp_size, the first
    (num_layers % pp_size) PP ranks get one extra layer. This exercises
    the layer distribution logic that mirrors the C++ getLayerNumPPRank
    pattern in cacheTransceiverTest.cpp.
    """
    print(
        f"\nRunning uneven PP transfer test: "
        f"ctx_tp={ctx_tp} ctx_pp={ctx_pp} gen_tp={gen_tp} gen_pp={gen_pp} "
        f"ctx_dp={ctx_enable_dp} gen_dp={gen_enable_dp} "
        f"mla={is_mla} v2={use_v2} num_layers={num_layers} "
        f"window={max_attention_window_vec} "
        f"layers_per_pp(ctx)={_get_layers_per_pp(num_layers, ctx_pp)} "
        f"layers_per_pp(gen)={_get_layers_per_pp(num_layers, gen_pp)}"
    )

    run_transfer_test(
        ctx_tp=ctx_tp,
        ctx_pp=ctx_pp,
        gen_tp=gen_tp,
        gen_pp=gen_pp,
        ctx_enable_dp=ctx_enable_dp,
        gen_enable_dp=gen_enable_dp,
        is_mla=is_mla,
        use_v2=use_v2,
        max_attention_window_vec=max_attention_window_vec,
        num_layers=num_layers,
    )

    print("PASSED")


@pytest.mark.timeout(180)
@pytest.mark.parametrize("is_mla", [False, True], ids=["mha", "mla"])
@pytest.mark.parametrize("gen_cp", [2, 4], ids=["cp2", "cp4"])
def test_cache_transceiver_helix_cp(gen_cp, is_mla):
    """Helix-CP KV transfer: un-sharded context (cp_size=1) -> helix generation.

    The context worker holds the whole global prompt; each generation CP rank
    owns the global blocks strided by [cp_rank::cp_size] (the ledger model, see
    transfer.py:1036 and KVCacheManagerV2._helix_local_len). verify_all_requests_helix
    asserts each gen rank received exactly its stride of the context blocks.

    Scope (deliberately narrow):
      * V2 only -- helix awareness lives in KVCacheManagerV2.
      * Full attention only -- the transfer path rejects a token_range (chunked /
        windowed slice) under CP (transfer.py:979), so no sliding window here.
      * tp=pp=1 on both sides -- keeps a rank == a CP rank so the verify is a
        plain positional stride without also unfolding TP head-concat.
      * No speculation -- helix rejects draft tokens.
    """
    print(
        f"\nRunning helix-CP transfer test: gen_cp={gen_cp} mla={is_mla} "
        f"(ctx cp=1, full attention, v2)"
    )

    run_transfer_test(
        ctx_tp=1,
        ctx_pp=1,
        gen_tp=1,
        gen_pp=1,
        ctx_enable_dp=False,
        gen_enable_dp=False,
        is_mla=is_mla,
        use_v2=True,
        max_attention_window_vec=None,
        gen_cp=gen_cp,
    )

    print("PASSED")


@pytest.mark.timeout(180)
@pytest.mark.parametrize("gen_cp", [2, 4, 8], ids=["cp2", "cp4", "cp8"])
def test_cache_transceiver_helix_cp_with_draft_layers(gen_cp):
    """Unified draft KV cache under helix: draft pages move like the target's.

    K-only draft layers after the MLA target's on both sides transfer with the
    target's block stride, so each generation CP rank holds the drafter's
    context pages [cp_rank::cp_size].
    """
    run_transfer_test(
        ctx_tp=1,
        ctx_pp=1,
        gen_tp=1,
        gen_pp=1,
        ctx_enable_dp=False,
        gen_enable_dp=False,
        is_mla=True,
        use_v2=True,
        max_attention_window_vec=None,
        gen_cp=gen_cp,
        num_draft_layers=2,
    )


# ---------------------------------------------------------------------------
# Boundary request-length tests
# ---------------------------------------------------------------------------
# Lengths chosen to exercise block-boundary and window-boundary edge cases:
#   TOKENS_PER_BLOCK=8, window_size=24
#   1         : single token (1 partial block)
#   7, 8, 9   : tpb-1, tpb, tpb+1 (block boundary)
#   23, 24, 25: win-1, win, win+1 (window boundary)
#   32, 33    : first stale_end=1 transition for win=24
BOUNDARY_REQUEST_LENGTHS = [1, 7, 8, 9, 23, 24, 25, 32, 33]

# Use a focused set of topologies to keep test count reasonable.
BOUNDARY_CONFIGS = [
    # (ctx_tp, ctx_pp, gen_tp, gen_pp, ctx_dp, gen_dp, is_mla, test_id)
    (1, 1, 1, 1, False, False, False, "tp1_pp1"),
    (2, 1, 2, 1, False, False, False, "tp2_pp1"),
    (1, 2, 1, 2, False, False, False, "tp1_pp2"),
    (2, 1, 1, 2, False, False, False, "tp2_to_pp2"),
    (1, 1, 1, 1, False, False, True, "mla_tp1_pp1"),
]


@pytest.mark.timeout(60)
@pytest.mark.parametrize("is_mla", [False, True], ids=["mha", "mla"])
def test_kv_cache_manager_v2_accepts_attention_dp_with_helix_cp(is_mla):
    """KVCacheManagerV2 builds a helix-CP generation instance with attention-DP at pp_size=1.

    The Python transceiver runs its gen-side completion consensus across the CP
    ranks of a DP group, and the pp=1 executor loops align gen-init admission
    across them.
    """
    managers = _create_managers_for_instance(
        tp=2,
        pp=1,
        enable_dp=True,
        is_mla=is_mla,
        use_v2=True,
        cp_size=2,
    )
    assert len(managers) == 4
    assert all(isinstance(m, KVCacheManagerV2) for m in managers)
    for m in managers:
        m.shutdown()


@pytest.mark.timeout(60)
def test_kv_cache_manager_v2_rejects_attention_dp_with_helix_cp_under_pp():
    """KVCacheManagerV2 rejects attention-DP with helix CP when pp_size > 1.

    Only the pp=1 executor loops align gen-init admission across helix CP ranks.
    """
    with pytest.raises(ValueError, match=r"pp_size > 1"):
        _create_managers_for_instance(
            tp=2,
            pp=2,
            enable_dp=True,
            is_mla=False,
            use_v2=True,
            cp_size=2,
        )


@pytest.mark.timeout(180)
@pytest.mark.parametrize(
    "ctx_tp,ctx_pp,gen_tp,gen_pp,ctx_enable_dp,gen_enable_dp,is_mla",
    [(c[0], c[1], c[2], c[3], c[4], c[5], c[6]) for c in BOUNDARY_CONFIGS],
    ids=[c[7] for c in BOUNDARY_CONFIGS],
)
@pytest.mark.parametrize(
    "max_attention_window_vec",
    [w[0] for w in WINDOW_CONFIGS],
    ids=[w[1] for w in WINDOW_CONFIGS],
)
@pytest.mark.parametrize("use_v2", [True], ids=["v2"])
def test_cache_transceiver_boundary_lengths(
    ctx_tp,
    ctx_pp,
    gen_tp,
    gen_pp,
    ctx_enable_dp,
    gen_enable_dp,
    is_mla,
    max_attention_window_vec,
    use_v2,
):
    """Test KvCacheTransceiverV2 with boundary request lengths.

    Exercises block-boundary and window-boundary edge cases
    (e.g. prompt_len == tokens_per_block, prompt_len == window_size ± 1).
    """
    print(
        f"\nRunning boundary transfer test: "
        f"ctx_tp={ctx_tp} ctx_pp={ctx_pp} gen_tp={gen_tp} gen_pp={gen_pp} "
        f"ctx_dp={ctx_enable_dp} gen_dp={gen_enable_dp} "
        f"mla={is_mla} v2={use_v2} window={max_attention_window_vec}"
    )

    run_transfer_test(
        ctx_tp=ctx_tp,
        ctx_pp=ctx_pp,
        gen_tp=gen_tp,
        gen_pp=gen_pp,
        ctx_enable_dp=ctx_enable_dp,
        gen_enable_dp=gen_enable_dp,
        is_mla=is_mla,
        use_v2=use_v2,
        max_attention_window_vec=max_attention_window_vec,
        request_lengths=BOUNDARY_REQUEST_LENGTHS,
    )

    print("PASSED")


@pytest.mark.timeout(180)
@pytest.mark.parametrize(
    (
        "ctx_tp",
        "ctx_pp",
        "gen_tp",
        "gen_pp",
        "is_mla",
        "use_v2",
        "pipelined_piece_tokens",
    ),
    [
        pytest.param(1, 1, 1, 1, False, True, None, id="v2_mha"),
        pytest.param(2, 1, 2, 1, True, True, None, id="v2_mla_tp2"),
        pytest.param(2, 1, 1, 2, False, True, None, id="v2_ctx_tp2_gen_pp2"),
        pytest.param(1, 1, 1, 1, False, True, 2 * TOKENS_PER_BLOCK, id="v2_mha_pipelined"),
    ],
)
def test_python_nixl_cache_transceiver_uses_cpp_bounce(
    ctx_tp: int,
    ctx_pp: int,
    gen_tp: int,
    gen_pp: int,
    is_mla: bool,
    use_v2: bool,
    pipelined_piece_tokens: Optional[int],
) -> None:
    """Exercise C++ bounce v2 through representative Python NIXL transceiver paths.

    Bounce is enabled by kv_cache_bounce_size_mb and tuned by agent_bounce_params via
    CacheTransceiverConfig inside run_transfer_test, which covers the config plumbing end to
    end and asserts engagement (agent.bounce_enabled / bounce_submit_count).
    """
    # BindingsNixlTransferStatus.last_status_str() resolves this exact attribute on the C++
    # status; if the binding name drifts, failure details silently degrade to "<unavailable>".
    from tensorrt_llm.tensorrt_llm_transfer_agent_binding import TransferStatus

    assert hasattr(TransferStatus, "get_last_status_str")

    run_transfer_test(
        ctx_tp=ctx_tp,
        ctx_pp=ctx_pp,
        gen_tp=gen_tp,
        gen_pp=gen_pp,
        ctx_enable_dp=False,
        gen_enable_dp=False,
        is_mla=is_mla,
        use_v2=use_v2,
        request_lengths=[30, 60],
        expect_cpp_bounce=True,
        pipelined_piece_tokens=pipelined_piece_tokens,
    )