Diff Coverage

    if query.shape[0] != 1:
        raise ValueError("Packed Lightning Indexer requires batch size one")
    plan = csa_tnd_plan(segments, ratio)
    if plan is None:
        return torch.full((*query.shape[:2], topk), -1, device=query.device, dtype=torch.int32)
    options, metadata = _tnd_indexer_metadata(plan, query.device, query.shape[2], query.shape[3], topk, ratio)
    with torch.autocast(device_type=query.device.type, enabled=False):
        indices, _ = _load_indexer_op()(
            plan.query(query), plan.key(key).unsqueeze(1), plan.query(weights).float(), topk,

    indices = sort_key_indices(indices, key.shape[1])
    indices = indices.masked_fill(indices == key.shape[1], -1)
    if plan.covers_queries(query.shape[1]):
        return indices.unsqueeze(0)
    result = torch.full((*query.shape[:2], topk), -1, device=query.device, dtype=torch.int32)
    plan.restore(indices, result)
    return result

    validate_indexer_tensors(query, key, weights)
    if indices.ndim != 3 or indices.shape[:2] != query.shape[:2]:
        raise ValueError("Fused KL indices must have shape [B, Q, TopK] matching Indexer Q")
    if key.shape[1] == 0 or not min(512, key.shape[1]) <= indices.shape[-1] <= 512:
        raise ValueError(f"Fused KL requires ordinary TopK512 selections, allowing padding for a compact bank; "
                         f"got K={key.shape[1]}, TopK={indices.shape[-1]}")
    if indices.device != query.device:
        raise ValueError(f"Fused KL indices must share Indexer Q's device, got {indices.device} and {query.device}")
    if query.device.type != "npu":

    handle: SequenceGatherHandle | None = None

    def start(self) -> SequenceGatherHandle:
        """Launch once in the differentiable forward, then release captured inputs."""
        if self.handle is None:
            self.handle = self.launcher()
            self.launcher = None
        return self.handle

    def wait(self) -> torch.Tensor:
        """Resolve the same differentiable handle at its first consumer."""
        return self.start().wait()


@dataclass
class SharedCompressedAttentionState:

        """
        if ratio <= 1 or (self.local_query_length % ratio == 0
                          and all(boundary % ratio == 0 for boundary in self.document_boundaries())):
            return None
        return document_compression_plan(self.document_boundaries(), ratio, self.local_query_length, size, rank)

    def compressed_positions(self, plan: Any, embeddings: tuple, context: Any) -> tuple:
        """Reuse boundary RoPE once per batch/source embedding, retaining its original precision.

    ) -> torch.Tensor:
        """Project queries with their original rotary positions, including moved tasks."""
        query = self.q_b_proj(query_residual.detach())
        if query.shape[-1] % self.head_dim:
            raise ValueError(
                f"Indexer query width {query.shape[-1]} is not divisible by head_dim {self.head_dim}"
            )
        query = query.view(*query_residual.shape[:2], -1, self.head_dim)
        if sequence_major:
            return _apply_v41_rope(query, cos, sin, unsqueeze_dim=2)
        return _apply_v41_rope(query.transpose(1, 2), cos, sin).transpose(1, 2)

    @torch.no_grad()
    def _project_worker_index_query(

        Returns:
            Worker Index Q in sequence-major layout, excluding padded rows.
        """
        if owner_query_length <= 0:
            raise ValueError("Indexer owner query length must be positive")
        batch, length = query_residual.shape[:2]
        if length == owner_query_length:
            return self._project_index_query(query_residual, cos, sin, sequence_major=True)
        if not length:
            query = self.q_b_proj(query_residual.detach())
            return query.view(batch, 0, query.shape[-1] // self.head_dim, self.head_dim)

        query = None
        for start in range(0, length, owner_query_length):
            end = min(start + owner_query_length, length)
            chunk = query_residual[:, start:end].detach()
            if end - start < owner_query_length:
                chunk = torch.nn.functional.pad(chunk, (0, 0, 0, owner_query_length - (end - start)))
            projected = self.q_b_proj(chunk)
            if query is None:
                query = projected.new_empty((batch, length, projected.shape[-1]))
            query[:, start:end].copy_(projected[:, :end - start])
            del projected, chunk
        query = query.view(batch, length, query.shape[-1] // self.head_dim, self.head_dim)
        return _apply_v41_rope(query, cos, sin, unsqueeze_dim=2)

    def _project_index_key(
            self,
            latent: torch.Tensor | None,

        keeps saving the same tensors in both passes.
        """
        cos, sin = compress_position_embeddings
        if cp_context is not None and cp_context.launch_balance is not None:
            return self._forward_dispatched(
                hidden_states, query_residual, latent, (cos, sin), index_key, cp_context,
                tp_context, document_lengths, prepare_attention, block_plan, key_position_embeddings,
            )
        key, key_handle = self._project_index_key(

        key, key_handle = self._project_index_key(
            latent, cos, sin, index_key, cp_context, block_plan, key_position_embeddings,
        )
        if prepare_attention is not None:
            prepare_attention()

        query = self._project_index_query(query_residual, cos, sin)
        local_num_heads = query.shape[2]
        merge_weight = self.weights_proj(hidden_states.detach())

            prepare_attention: Callable[[], None] | None,
            block_plan: Any = None, key_position_embeddings: tuple | None = None,
    ) -> SharedCompressedIndexerOutput:
        """Overlap worker inputs and key gathering while retaining owner-side KL edges."""
        cos, sin = embeddings
        if tp_context is not None or self.is_candidate_source or self.uses_candidates:
            raise NotImplementedError("Indexer overlap requires ordinary causal selection and TP=1")
        weights = self.weights_proj(hidden_states.detach())
        weights = weights * (self.head_dim**-0.5 * self.global_num_heads**-0.5)
        launch = self._launch_selection if self._input_region is None else self._input_region
        wave = launch(query_residual, weights, (cos, sin), document_lengths, context)
        key, handle = self._project_index_key(
            latent, cos, sin, index_key, context, block_plan, key_position_embeddings,
        )
        if prepare_attention is not None:
            prepare_attention()
        prepare = self._prepare_dispatched if self._prepare_region is None else self._prepare_region
        prepared = prepare(wave) if context.indexer_priority else None
        # KL uses this autograd edge; only selection transport runs under no_grad.
        if handle is not None:
            key = handle.wait()
        if not context.indexer_priority:
            prepared = prepare(wave)
        select = self._select_dispatched if self._result_region is None else self._result_region
        result = select(prepared, key)
        del prepared
        query = self._project_index_query(query_residual, cos, sin)
        if weights.shape[-1] != query.shape[2]:
            raise ValueError("Indexer query and merge weights must have matching head counts")
        finish = self._finish_selection if self._wait_region is None else self._wait_region
        return SharedCompressedIndexerOutput(finish(result), key, None, query, weights)

    def enable_selection_exclude(self) -> None:
        """Keep discrete selection fixed when the enclosing layer recomputes."""
        object.__setattr__(

        """Keep discrete selection fixed when the enclosing layer recomputes."""
        object.__setattr__(
            self, "_selection_region", checkpoint_exclude_wrapper(self._select_tokens)
        )
        object.__setattr__(self, "_input_region", checkpoint_exclude_wrapper(self._launch_selection))
        # This output goes straight into an exclusion region; retain IDs, not expanded worker Q.
        object.__setattr__(self, "_prepare_region", checkpoint_exclude_wrapper(
            self._prepare_dispatched, save_output=False,
        ))
        object.__setattr__(self, "_result_region", checkpoint_exclude_wrapper(self._select_dispatched))
        object.__setattr__(self, "_wait_region", checkpoint_exclude_wrapper(self._finish_selection))

    @torch.no_grad()
    def _launch_selection(self, qr: torch.Tensor, weights: torch.Tensor, embeddings: tuple,
                          lengths: tuple, context: SharedCompressedAttentionCPContext) -> Any:

    @torch.no_grad()
    def _launch_selection(self, qr: torch.Tensor, weights: torch.Tensor, embeddings: tuple,
                          lengths: tuple, context: SharedCompressedAttentionCPContext) -> Any:
        """Cache the input launch so checkpoint replay does not repeat communication."""
        return context.launch_balance(qr, weights, embeddings, lengths, self.compress_ratio)

    @torch.no_grad()
    def _prepare_dispatched(self, wave: Any) -> tuple:
        """Expand received qr while Index K may still be in flight."""
        qr, weights, cos, sin = wave.wait()
        query = self._project_worker_index_query(qr, cos, sin, wave.owner_queries)
        return query, weights, wave

    @torch.no_grad()
    def _select_dispatched(self, prepared: tuple, key: torch.Tensor) -> Any:
        """Consume worker Q once and launch the inverse ID transfer."""
        query, weights, wave = prepared
        selected = self._select_indices(query, key, weights, None, 0, None, None, wave.route.segments)[0]
        return wave.launch_result(selected[0])

    @staticmethod
    @torch.no_grad()
    def _finish_selection(result: Any) -> torch.Tensor:

    @staticmethod
    @torch.no_grad()
    def _finish_selection(result: Any) -> torch.Tensor:
        """Cache completed IDs; replay never consumes an already finished handle."""
        return result.wait().unsqueeze(0)

    def _run_selection(self, *args: Any, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Select directly, or reuse the result retained by the checkpoint region."""
        selection = self._select_tokens if self._selection_region is None else self._selection_region

    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """Cache dispatch, discrete selection and TopK return as one recompute region."""
        cos, sin = embeddings
        if cp_context is not None and cp_context.balance_indexer is not None:
            if tp_context is not None or self.is_candidate_source or self.uses_candidates:
                raise NotImplementedError("Indexer load balance currently requires ordinary causal selection and TP=1")

            def select_tasks(task_qr: torch.Tensor, task_weights: torch.Tensor,
                             task_embeddings: tuple, segments: tuple) -> torch.Tensor:
                """Project on the receiver while this layer's parameters are resident."""
                task_query = self._project_worker_index_query(
                    task_qr, *task_embeddings, query_residual.shape[1],
                )
                return self._select_indices(
                    task_query, key, task_weights, None, 0, None, None, segments,
                )[0]

            topk_indices = cp_context.balance_indexer(
                query_residual, merge_weight, (cos, sin), select_tasks, document_lengths, self.compress_ratio,
            )
            candidate_blocks = None
        else:
            topk_indices, candidate_blocks = self._select_indices(
                query, key, merge_weight, candidate_blocks, query_offset,
                reduce_sum, minimum_key_indices, query_segments,

        key_value = _apply_v41_rope(key_value, cos, sin)
        if cp_context is None:
            return key_value, None
        if raw_route is not None:
            return None, cp_context.launch_raw(key_value, 2, raw_route)
        return None, cp_context.launch(key_value, 2)

    def _project_query(
            self,

        if packed_sequence is not None:
            boundaries = packed_sequence.prepare(device, (self.compress_ratio,)).document_boundaries()
        raw_route, main_route = None, None
        if cp_context is not None and cp_context.launch_raw is not None:
            raw_route = sliding_window_route(
                boundaries, sequence_length, cp_context.size, cp_context.rank, self.sliding_window,
            )
        if self.compress_ratio and cp_context is not None and cp_context.launch_main is not None:
            main_route = document_prefix_route(
                boundaries, sequence_length, cp_context.size, cp_context.rank, self.compress_ratio,
            )
        return raw_route, main_route, tuple(end - start for start, end in zip(boundaries, boundaries[1:]))

    def _launch_compressed_bank(compressed: torch.Tensor, context: SharedCompressedAttentionCPContext,
                                route: SequenceDispatchRoute | None, block_plan: Any) -> SequenceGatherHandle:
        """Use the configured physical bank layout at the scheduled launch point."""
        if route is not None:
            return context.launch_main(compressed, 1, route)
        if block_plan is not None:
            return block_plan.launch_bank(compressed, context)
        return context.launch(compressed, 1)

    def _run_indexer(
            self,

                shared_state.publish_compressed_kv(
                    self.layer_idx, compressed, 0 if main_route is None else main_route.global_start,
                )
        if key_value is None:
            raise RuntimeError("raw KV communication did not produce a tensor")

        return key_value

    def _build_sparse_indices(

                and compressed_kv is not None
        ):
            # KL uses the original owner's queries. A view aligns its student K
            # with the compact teacher bank while preserving global Index K gradients.
            index_key = indexer_output.index_key
            if main_key_offset or compressed_kv.shape[1] != index_key.shape[1]:
                index_key = index_key[:, main_key_offset:main_key_offset+compressed_kv.shape[1]]
            selected = indexer_output.topk_indices
            if main_key_offset:
                selected = torch.where(selected >= 0, selected - main_key_offset, -1)
                query_segments = tuple((s, e, ks-main_key_offset, ke-main_key_offset, residual)
                                       for s, e, ks, ke, residual in query_segments)
            indexer_loss = shared_compressed_indexer_kl_loss(
                indexer_output.index_query,
                index_key,

        overlap = (self.is_index_source and cp_context is not None
                   and (cp_context.indexer_priority or cp_context.launch_balance is not None))
        prepare_attention = None
        if overlap:
            query_residual = self.q_a_norm(self.q_a_proj(hidden_states))
            query = None

            def _prepare_attention() -> None:
                """Keep differentiable main Q work outside selection checkpoint regions."""
                nonlocal query
                if isinstance(compressed_handle, _DeferredBankLaunch):
                    compressed_handle.start()
                query = self._expand_query(query_residual, cos, sin)

            prepare_attention = _prepare_attention
        else:
            query_residual, query = self._project_query(hidden_states, cos, sin)
        latent, compressed_handle = self._compute_compressed_kv(
            hidden_states, position_embeddings, shared_state, cp_context, main_route, block_plan, key_positions,

        ctx: Any,
        grad_output: Tensor,
    ) -> tuple[Tensor, None, None, None, None, None]:
        """SUM consumer gradients directly into each owner's sequence shard."""
        send = _move_dim_to_front(grad_output, ctx.gather_dim)
        output = grad_output.new_empty((send.shape[0] // ctx.world_size, *send.shape[1:]))
        work = dist.reduce_scatter_tensor(
            output, send, op=dist.ReduceOp.SUM, group=ctx.group, async_op=True,
        )
        work.wait()
        local = _move_dim_from_front(output, ctx.gather_dim)
        return (
            local,
            None,
            None,

    None represents an omitted option so none/zigzag remain valid without it.
    Explicit values require greedy/hybrid mode and must be positive integers.
    """
    if mode not in ("none", "zigzag", "greedy", "hybrid"):
        raise ValueError("indexer_load_balance must be 'none', 'zigzag', 'greedy' or 'hybrid'")
    if block_size is not None:
        if mode not in ("greedy", "hybrid"):
            raise ValueError("indexer_balance_block_size requires indexer_load_balance='greedy' or 'hybrid'")
        if not isinstance(block_size, int) or isinstance(block_size, bool) or block_size < 1:
            raise ValueError("indexer_balance_block_size must be a positive integer")
    return (2048 if block_size is None else block_size) if mode in ("greedy", "hybrid") else None


def validate_indexer_hybrid_threshold(mode: str, min_length: int | None) -> None:

        min_length: Explicit positive raw-token document threshold for hybrid.
    """
    if mode != "hybrid":
        if min_length is not None:
            raise ValueError("indexer_zigzag_min_length requires indexer_load_balance='hybrid'")
    elif not isinstance(min_length, int) or isinstance(min_length, bool) or min_length < 1:
        raise ValueError(
            "indexer_load_balance='hybrid' requires an explicit positive integer indexer_zigzag_min_length"
        )

        for _ in range(len(available)):
            cost, count, rank = heappop(available)
            if count + task.length <= capacity:
                break
            skipped.append((cost, count, rank))
        else:
            raise ValueError(f"No Indexer worker has capacity for {task.length} query rows (limit: {capacity})")
        assignments[rank].append(task)
        heappush(available, (cost - negative_work, count + task.length, rank))
        for item in skipped:
            heappush(available, item)


def _hybrid_tasks(lengths: tuple[int, ...], size: int, ratio: int, block: int,
                   min_length: int) -> list[list[IndexerTask]]:

def _validate_task_plan(lengths: tuple[int, ...], size: int, ratio: int, mode: str,
                        block_size: int | None) -> int | None:
    """Keep input validation identical for direct and all-worker planning."""
    if size < 1 or ratio < 1 or not lengths or min(lengths) < 1 or sum(lengths) % size:
        raise ValueError("Indexer balance needs positive document lengths/ratio and an equal contiguous CP partition")
    block = resolve_indexer_balance_block_size(mode, block_size)
    if mode not in ("zigzag", "greedy", "hybrid"):
        raise ValueError("Indexer task mode must be 'zigzag', 'greedy' or 'hybrid'")
    local = sum(lengths) // size
    if block is not None and block > local:
        raise ValueError(f"Indexer greedy block size {block} exceeds the local query count {local}")
    return block


@lru_cache(maxsize=128)

    Returns:
        Immutable input/result route with a local-task bypass.
    """
    if not 0 <= rank < size:
        raise ValueError("Indexer balance rank must be inside its CP group")
    if mode == "zigzag":
        _validate_task_plan(lengths, size, ratio, mode, block_size)
        validate_indexer_hybrid_threshold(mode, zigzag_min_length)
        boundaries = (0, *accumulate(lengths))

@lru_cache(maxsize=32)
@torch.inference_mode(False)
def _device_rows(route: IndexerRoute, device: torch.device) -> tuple:
    return tuple(torch.tensor(rows, device=device, dtype=torch.long) for rows in
                 (route.send_rows, route.receive_rows, route.local_sources, route.local_targets))


def _exchange_rows(tensor: torch.Tensor, route: IndexerRoute, group: Any, rows: tuple) -> torch.Tensor:

def _exchange_rows(tensor: torch.Tensor, route: IndexerRoute, group: Any, rows: tuple) -> torch.Tensor:
    """Exchange remote rows and place local bypass rows in worker task order."""
    send_rows, receive_rows, local_sources, local_targets = rows
    send = tensor.index_select(0, send_rows).contiguous()
    received = tensor.new_empty((sum(route.output_splits), *tensor.shape[1:]))
    work = dist.all_to_all_single(received, send, list(route.output_splits), list(route.input_splits),
                                  group=group, async_op=True)
    output = tensor.new_empty((route.query_count, *tensor.shape[1:]))
    output.index_copy_(0, local_targets, tensor.index_select(0, local_sources))
    work.wait()
    output.index_copy_(0, receive_rows, received)
    return output


def _exchange_inputs(inputs: tuple[torch.Tensor, ...], route: IndexerRoute, group: Any, rows: tuple) -> tuple:
    """Pack fields of equal dtype together without rounding the original RoPE values."""
    by_dtype = {}
    for index, tensor in enumerate(inputs):
        by_dtype.setdefault(tensor.dtype, []).append(index)
    outputs = [None] * len(inputs)
    for indices in by_dtype.values():
        widths = [inputs[index].shape[-1] for index in indices]
        packed = torch.cat([inputs[index][0] for index in indices], dim=-1)
        received = _exchange_rows(packed, route, group, rows)
        for index, field in zip(indices, received.split(widths, dim=-1)):
            outputs[index] = field.unsqueeze(0).contiguous()
    return tuple(outputs)


@torch.no_grad()
def balanced_indexer_select(query_residual: torch.Tensor, weights: torch.Tensor,

        group: Original sequence-ordered CP process group.
        block_size: Greedy task length in original query rows; see indexer_task_plan.
        zigzag_min_length: Hybrid-only minimum original document length.
    """
    size, rank = dist.get_world_size(group), dist.get_rank(group)
    inputs = (query_residual, weights, *position_embeddings)
    if (len(position_embeddings) != 2 or query_residual.ndim != 3 or query_residual.shape[0] != 1
            or query_residual.shape[1] * size != sum(lengths)):
        raise ValueError("Indexer balance requires batch one and equal contiguous source query shards")
    if any(tensor.ndim != 3 or tensor.shape[:2] != query_residual.shape[:2]
           or tensor.device != query_residual.device for tensor in inputs):
        raise ValueError("Indexer qr, weights and rotary embeddings must have matching batch/token axes and device")
    route = indexer_route(lengths, size, rank, ratio, mode, block_size,
                           zigzag_min_length)
    rows = _device_rows(route, query_residual.device)
    balanced_qr, balanced_weights, cos, sin = _exchange_inputs(inputs, route, group, rows)
    selected = selector(balanced_qr, balanced_weights, (cos, sin), route.segments)[0]
    send_rows, receive_rows, local_sources, local_targets = rows
    send = selected.index_select(0, receive_rows).contiguous()
    received = selected.new_empty((sum(route.input_splits), selected.shape[-1]))
    work = dist.all_to_all_single(received, send, list(route.input_splits), list(route.output_splits),
                                  group=group, async_op=True)
    output = selected.new_empty((query_residual.shape[1], selected.shape[-1]))
    output.index_copy_(0, local_sources, selected.index_select(0, local_targets))
    work.wait()
    output.index_copy_(0, send_rows, received)
    return output.unsqueeze(0)


@dataclass
class _PendingRows:

    @torch.no_grad()
    def wait(self) -> torch.Tensor:
        """Wait once, then restore received rows and release transfer storage."""
        if self.consumed:
            raise RuntimeError("Indexer communication handle cannot be consumed twice")
        self.consumed = True
        self.work.wait()
        output = self.output
        output.index_copy_(0, self.receive_rows, self.received)
        self.send = self.received = self.output = self.work = None
        return output


@dataclass
class IndexerInputExchange:

    @torch.no_grad()
    def wait(self) -> tuple:
        """Restore qr, weights and original RoPE after all input launches."""
        if self.consumed:
            raise RuntimeError("Indexer input exchange cannot be consumed twice")
        self.consumed = True
        outputs = [None] * len(self.widths)
        for transfer, indices in zip(self.transfers, self.fields):
            merged = transfer.wait()
            for index, value in zip(indices, merged.split([self.widths[i] for i in indices], dim=-1)):
                outputs[index] = value.unsqueeze(0).contiguous()
        self.transfers = ()
        return tuple(outputs)

    @torch.no_grad()
    def launch_result(self, selected: torch.Tensor) -> _PendingRows:
        """Return selection IDs while owner-side differentiable projections run."""
        send_rows, receive_rows, local_sources, local_targets = self.rows
        send = selected.index_select(0, receive_rows).contiguous()
        received = selected.new_empty((sum(self.route.input_splits), selected.shape[-1]))
        work = dist.all_to_all_single(
            received, send, list(self.route.input_splits), list(self.route.output_splits),
            group=self.group, async_op=True,
        )
        output = selected.new_empty((self.owner_queries, selected.shape[-1]))
        output.index_copy_(0, local_sources, selected.index_select(0, local_targets))
        return _PendingRows(send, received, output, send_rows, work)


@torch.no_grad()
def launch_indexer_inputs(

    Returns:
        Single-use handle retaining transfer storage until wait.
    """
    size, rank = dist.get_world_size(group), dist.get_rank(group)
    inputs = (query_residual, weights, *embeddings)
    if (len(embeddings) != 2 or query_residual.ndim != 3 or query_residual.shape[0] != 1
            or query_residual.shape[1] * size != sum(lengths)):
        raise ValueError("Indexer balance requires batch one and equal contiguous source query shards")
    if any(value.ndim != 3 or value.shape[:2] != query_residual.shape[:2]
           or value.device != query_residual.device for value in inputs):
        raise ValueError("Indexer qr, weights and rotary embeddings must share token axes and device")
    route = indexer_route(lengths, size, rank, ratio, mode, block_size,
                          zigzag_min_length)
    rows = _device_rows(route, query_residual.device)
    send_rows, receive_rows, local_sources, local_targets = rows
    grouped = {}
    for index, value in enumerate(inputs):
        grouped.setdefault(value.dtype, []).append(index)
    transfers, fields = [], []
    for indices in grouped.values():
        packed = torch.cat([inputs[index][0] for index in indices], dim=-1)
        send = packed.index_select(0, send_rows).contiguous()
        received = packed.new_empty((sum(route.output_splits), packed.shape[-1]))
        work = dist.all_to_all_single(
            received, send, list(route.output_splits), list(route.input_splits), group=group, async_op=True,
        )
        output = packed.new_empty((route.query_count, packed.shape[-1]))
        output.index_copy_(0, local_targets, packed.index_select(0, local_sources))
        transfers.append(_PendingRows(send, received, output, receive_rows, work))
        fields.append(tuple(indices))
    return IndexerInputExchange(route, rows, group, query_residual.shape[1], tuple(transfers),
                                tuple(fields), tuple(value.shape[-1] for value in inputs))

        route: Immutable per-batch routing agreed by every member of the CP group.
        group: Existing framework-owned CP process group, in sequence-shard rank order.
    """
    if not -tensor.ndim <= sequence_dim < tensor.ndim:
        raise ValueError("sequence dispatch dimension is outside the input tensor dimensions")
    sequence_dim %= tensor.ndim
    if tensor.shape[sequence_dim] != route.local_length:
        raise ValueError("sequence dispatch owner length does not match the bank's physical layout")
    local = tensor.movedim(sequence_dim, 0).contiguous()

def sliding_window_route(boundaries: tuple[int, ...], local_length: int, size: int,
                         rank: int, window: int) -> SequenceDispatchRoute:
    """Plan the same-document left halo; every rank agrees whether to exchange."""
    if local_length <= 0 or size <= 0 or not 0 <= rank < size or window <= 0:
        raise ValueError("SWA dispatch requires positive dimensions and a valid CP rank")
    if (len(boundaries) < 2 or boundaries[0] != 0 or boundaries[-1] != local_length * size
            or any(end <= start for start, end in zip(boundaries, boundaries[1:]))):
        raise ValueError("SWA document boundaries must span the full contiguous CP sequence")
    intervals = tuple(
        (max(boundaries[bisect_right(boundaries, peer * local_length) - 1], peer * local_length - window + 1),
         (peer + 1) * local_length)
        for peer in range(size)

def document_prefix_route(boundaries: tuple[int, ...], local_length: int, size: int,
                          rank: int, ratio: int) -> SequenceDispatchRoute:
    """Route natural-document main KV prefixes using complete-block ownership."""
    if ratio not in (1, 2) or local_length <= 0 or size <= 0 or not 0 <= rank < size:
        raise ValueError("Main KV dispatch requires a supported ratio and equal positive raw CP shards")
    if (len(boundaries) < 2 or boundaries[0] != 0 or boundaries[-1] != local_length * size
            or any(end <= start for start, end in zip(boundaries, boundaries[1:]))):
        raise ValueError("Main KV document boundaries must span the complete contiguous CP sequence")
    keys = compressed_document_offsets(boundaries, ratio)
    owners = compressed_owner_offsets(boundaries, ratio, local_length)
    intervals = tuple(
        (keys[bisect_right(boundaries, peer * local_length) - 1], owners[peer + 1])