Diff Coverage

    Raises:
        ValueError: Tensor layouts, devices or query/key dtypes disagree.
    """
    if query.ndim != 4:
        raise ValueError(f"Indexer query must have shape [B, Q, H, D], got {tuple(query.shape)}")
    if key.ndim != 3 or key.shape[::2] != (query.shape[0], query.shape[-1]):
        raise ValueError(f"Indexer key must have shape [B, K, D] matching Q, got Q={tuple(query.shape)}, "
                         f"K={tuple(key.shape)}")
    if weights.shape != query.shape[:3]:
        raise ValueError(f"Indexer weights must have shape {tuple(query.shape[:3])}, got {tuple(weights.shape)}")
    if key.device != query.device or weights.device != query.device:
        raise ValueError(f"Indexer Q/K/weights must share a device, got {query.device}, {key.device}, {weights.device}")
    if key.dtype != query.dtype:
        raise ValueError(f"Indexer Q/K must share a dtype, got {query.dtype} and {key.dtype}")


def sort_key_indices(indices: torch.Tensor, maximum: int) -> torch.Tensor:
    """Sort key IDs without losing integer precision or changing their dtype.

            "use_fused_ops=True requires cann_ops_transformer.lightning_indexer_metadata on this device. "
            "See hyper_parallel/components/functional/README.md for ops-transformer installation, "
            "environment setup and tested versions."
        ) from err
    return lightning_indexer_metadata


@lru_cache(maxsize=16)
def _uses_indexer_metadata(device: torch.device) -> bool:

    LI V2's arch22 kernel computes its split internally; arch35 requires
    metadata. Keep the standard metadata call on other devices rather than
    assuming they share the verified 910B behavior. TND still needs cu_seqlens.
    """
    return "910B" not in torch.npu.get_device_name(device)


@lru_cache(maxsize=128)
@torch.inference_mode(False)

        topk: Output selection width, including invalid tail slots.
        segment: Local Q start/end, global compressed K start/end, and raw
            prefix length modulo ratio, shared by all physical batch items.
    """
    batch, length, heads, dim = query.shape
    result = torch.full((batch, length, topk), -1, device=query.device, dtype=torch.int32)
    start, end, key_start, key_end, residual = segment
    if key_end <= key_start or start == end:
        return result
    indexer = _load_indexer_op()
    residue, metadata = _bsnd_indexer_metadata(
        query.device, batch, end-start, heads, dim, key_end-key_start, topk, ratio, residual,
    )
    with torch.autocast(device_type=query.device.type, enabled=False):
        indices, _ = indexer(
            query[:, start:end].contiguous(), key[:, key_start:key_end].unsqueeze(2).contiguous(),
            weights[:, start:end].float().contiguous(), topk,
            cmp_residual_k=residue, metadata=metadata,
            layout_q="BSND", layout_k="BSND", mask_mode=3, cmp_ratio=ratio,

            weights[:, start:end].float().contiguous(), topk,
            cmp_residual_k=residue, metadata=metadata,
            layout_q="BSND", layout_k="BSND", mask_mode=3, cmp_ratio=ratio,
        )
    indices = indices[:, :, 0]
    indices = torch.where(indices >= 0, indices + key_start, key.shape[1])
    indices = sort_key_indices(indices, key.shape[1])
    result[:, start:end] = indices.masked_fill(indices == key.shape[1], -1)
    return result


@lru_cache(maxsize=128)
@torch.inference_mode(False)

    """Cache native selection metadata for the complete packed document partition."""
    options = plan.indexer_options(device, ratio)
    metadata = None
    if _uses_indexer_metadata(device):
        metadata = _load_indexer_metadata_op()(
            heads, 1, dim, topk, batch_size=len(plan.segments),
            max_seqlen_q=max(plan.q_lengths), max_seqlen_k=max(plan.k_lengths), **options,
        )
    return options, metadata

    Returns:
        Sorted global key IDs [batch, queries, topk], with invalid slots -1.
    """
    if len(segments) == 1:
        return _bsnd_lightning_indexer(query, key, weights, ratio, topk, segments[0])
    if query.shape[0] != 1:
        raise ValueError("Packed Lightning Indexer requires batch size one")
    result = torch.full((*query.shape[:2], topk), -1, device=query.device, dtype=torch.int32)
    plan = csa_tnd_plan(segments, ratio)
    if plan is None:
        return result
    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,

        selected = self.query(indices)
        offsets = self.lengths(indices.device)[3]
        selected = torch.where(selected >= 0, selected - offsets, -1).int()
        if selected.shape[-1] < padded_width:
            selected = torch.nn.functional.pad(selected, (0, padded_width-selected.shape[-1]), value=-1)
        return selected.unsqueeze(1).contiguous()

    def covers_queries(self, length: int) -> bool:
        """Return whether packed rows already follow the complete original Q order."""

            "and sparse_flash_mla_metadata. "
            "See hyper_parallel/components/functional/README.md for ops-transformer installation, "
            "environment setup and tested versions."
        ) from err
    return sparse_flash_mla_op, sparse_flash_mla_grad, sparse_flash_mla_metadata


def supports_fused_attention_inputs(query: torch.Tensor, ratio: int, window: int, topk: int = 512) -> bool:
    """Check the validated attention input contract without importing operators.

        topk: Logical compressed selection width, before padding.
    """
    if query.device.type != "npu" or query.dtype != torch.bfloat16:
        return False
    if query.shape[-1] != 512 or window != 128 or ratio not in (0, 1, 2):
        return False
    if query.shape[1] not in (1, 2, 4, 8, 16, 32, 64, 128) or query.shape[2] == 0:
        return False
    if ratio and not 1 <= topk <= 512:
        return False
    return True


@lru_cache(maxsize=128)
@torch.inference_mode(False)

    Returns:
        Attention output in [B, Q, H, D].
    """
    if query.dtype != torch.bfloat16 or query.shape[-1] != 512 or window != 128:
        raise ValueError("V4 attention requires BF16, head_dim=512, and window=128")
    start, end, key_start, key_end, residual = segment
    document_start = key_start * max(ratio, 1)
    raw_left = max(document_start, query_offset + start - window + 1)
    raw_right = query_offset + end

    @staticmethod
    def backward(ctx: Any, grad_output: torch.Tensor) -> tuple:
        """Return each physical bank's gradient to its autograd communication owner."""
        query, raw, main, indices, sinks, output, lse, cu_q, cu_raw, cu_k, residuals = ctx.saved_tensors
        _, backward, _ = _load_attention_ops()
        dq, dr, dm, ds, _, _ = backward(
            query, grad_output.contiguous(), output, lse, ori_kv=raw, cmp_kv=main,
            cmp_sparse_indices=indices, sinks=sinks, cu_seqlens_q=cu_q, cu_seqlens_ori_kv=cu_raw,
            cu_seqlens_cmp_kv=cu_k, cmp_residual_kv=residuals, **ctx.options,
        )
        return dq, dr, dm if main is not None else None, None, ds, None, None, None, None, None


def sparse_flash_mla(
        query: torch.Tensor, raw_key: torch.Tensor, main_key: torch.Tensor | None,

        return _bsnd_sparse_flash_mla(
            query, raw_key, main_key, indices, sinks, scale, ratio, window, query_offset, segments[0],
        )
    if query.shape[0] != 1:
        raise ValueError("Packed Sparse Flash MLA requires batch size one")
    if query.dtype != torch.bfloat16 or query.shape[-1] != 512 or window != 128:
        raise ValueError("V4 attention requires BF16, head_dim=512, and window=128")
    output = None
    for has_keys in (True, False):
        plan = csa_tnd_plan(segments, ratio, has_keys)
        if plan is None:
            continue
        raw_ranges = tuple(
            (max(ks * max(ratio, 1), query_offset+s-window+1), query_offset+e)
            for s, e, ks, _, _ in plan.segments
        )

        )
        # A complete Q partition is already ordered; a canvas would retain another full output.
        if plan.covers_queries(query.shape[2]):
            return part.unsqueeze(0)
        if output is None:
            output = torch.zeros_like(query.transpose(1, 2))
        plan.restore(part, output)
    return torch.zeros_like(query.transpose(1, 2)) if output is None else output

            "and sparse_lightning_indexer_kl_loss_grad_metadata. "
            "See hyper_parallel/components/functional/README.md for ops-transformer installation, "
            "environment setup and tested versions."
        ) from err
    return sparse_lightning_indexer_kl_loss_grad, sparse_lightning_indexer_kl_loss_grad_metadata


def validate_fused_kl_inputs(
        query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor, indices: torch.Tensor,

        RuntimeError: Explicit fusion was requested on a non-NPU device.
    """
    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 indices.shape[-1] != min(512, key.shape[1]):
        raise ValueError(f"Fused KL requires ordinary TopK512 selections from a nonempty 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":
        raise RuntimeError(f"Fused KL requires NPU inputs; got device={query.device}")

        device: torch.device, batch: int, length: int, heads: int, dim: int,
        key_length: int, topk: int, ratio: int, residual: int,
) -> tuple:
    """Cache shape-only causal metadata without retaining activations."""
    residue = torch.full((batch,), residual, device=device, dtype=torch.int32) if ratio > 1 else None
    options = {"layout_q": "BSND", "layout_k": "BSND", "mask_mode": 3,
               "cmp_ratio": ratio, "cmp_residual_k": residue}
    _, metadata_op = _load_kl_ops()
    metadata = metadata_op(
        heads, 1, dim, batch_size=batch, max_seqlen_q=length, max_seqlen_k=key_length,
        topk=topk, **options,
    )
    return options, metadata


def gather_selected_keys(bank: torch.Tensor, indices: torch.Tensor) -> torch.Tensor:
    """Gather batch-local nonnegative indices without multidimensional indexing.

    with torch.no_grad(), torch.autocast(device_type=query.device.type, enabled=False):
        batch, heads, length, dim = query.shape
        width = indices.shape[-1]
        if width == 0:
            return query.new_zeros(indices.shape, dtype=torch.float32)
        valid = indices >= 0
        safe = indices.clamp_min(0).long()
        selected_chunk = min(width, max(1, _FP32_TEMPORARY_BYTES // max(1, batch * length * dim * 4)))
        head_chunk = min(heads, max(1, _FP32_TEMPORARY_BYTES // max(1, batch * length * width * 4)))

            probabilities = torch.cat((scores, sink_logits), dim=-1).softmax(dim=-1)
            target.add_(probabilities[..., :-1].masked_fill(~valid.unsqueeze(1), 0).sum(dim=1))
            del scores, probabilities, queries
        if reduce_sum is not None:
            target = reduce_sum(target)
        return target / target.sum(dim=-1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny)


def causal_attention_teacher(

    log_prediction = prediction.masked_fill(~valid, 1).log()
    # The native interface does not return student LSE. One boundary check
    # avoids silently clamping the scalar objective when probabilities vanish.
    if not bool(bad.any()):
        return log_prediction
    for batch in range(query.shape[0]):
        rows = bad[batch].nonzero().flatten()
        for start in range(0, rows.numel(), chunk_size):
            chosen = rows[start:start+chunk_size]

        Returns:
            Scalar KL normalized by this document segment's batch and query count.
        """
        batch, length, heads, dim = query.shape
        kl_grad, _ = _load_kl_ops()
        with torch.no_grad(), torch.autocast(device_type=query.device.type, enabled=False):
            target = causal_attention_teacher(
                attention_query, main_key, indices, sinks, scale, ratio, residual, chunk_size,
            )
            options, metadata = _kl_metadata(
                query.device, batch, length, heads, dim, key.shape[1], indices.shape[-1], ratio, residual,
            )
            dq, dk, dw, prediction = kl_grad(
                query.contiguous(), key.unsqueeze(2).contiguous(), weights.float().contiguous(),
                indices.unsqueeze(2).int().contiguous(), target.unsqueeze(2), metadata=metadata, **options,
            )
            log_prediction = _stable_prediction_log(prediction[:, :, 0], query, key, weights, indices, chunk_size)
            loss = (target * (target.clamp_min(torch.finfo(torch.float32).tiny).log() - log_prediction)).sum()
            ctx.save_for_backward(dq, dk.squeeze(2), dw)
            ctx.factor = coefficient / (batch * length)
        return loss * ctx.factor

    @staticmethod
    def backward(ctx: Any, seed: torch.Tensor) -> tuple:
        """Apply the caller's seed in FP32, without differentiating the teacher.

        Returns:
            Student Q/K/weight gradients and None for teacher and sequence metadata.
        """
        gradients = tuple(value.float() * (seed * ctx.factor) for value in ctx.saved_tensors)
        return (*gradients, *((None,) * 9))


def _bsnd_sparse_lightning_indexer_kl_loss(
        query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor,

        segment: Local Q start/end, document K start/prefix end and residual,
            shared by all physical batch items.
    """
    # Zero-length views attach zero gradients even when the segment has no key.
    total = query[..., :0].float().sum() + key[..., :0].float().sum() + weights[..., :0].float().sum()
    start, end, key_start, key_end, residual = segment
    if key_end <= key_start or start == end:
        return total
    selected = indices[:, start:end]
    selected = torch.where(selected >= 0, selected - key_start, -1)
    if selected.shape[-1] < 512:
        selected = torch.nn.functional.pad(selected, (0, 512-selected.shape[-1]), value=-1)
    loss = _SparseLightningIndexerKLLoss.apply(
        query[:, start:end], key[:, key_start:key_end], weights[:, start:end],
        attention_query[:, :, start:end], main_key[:, key_start:key_end], selected, sinks,
        scale, coefficient, chunk_size, ratio, residual,
    )
    return total + loss * ((end-start) / query.shape[1])


@lru_cache(maxsize=128)
@torch.inference_mode(False)

@lru_cache(maxsize=128)
@torch.inference_mode(False)
def _tnd_kl_metadata(plan, device, heads, dim, ratio):
    """Cache KL metadata using the same Q/K partition as native selection."""
    options = plan.indexer_options(device, ratio)
    _, metadata_op = _load_kl_ops()
    metadata = metadata_op(
        heads, 1, dim, topk=512, batch_size=len(plan.segments),
        max_seqlen_q=max(plan.q_lengths), max_seqlen_k=max(plan.k_lengths), **options,
    )
    return options, metadata


class _SparseLightningIndexerKLLossTND(torch.autograd.Function):
    """Fuse the student gradients for several documents, retaining Hyper's teacher."""

            ctx: Any, query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor,
            indices: torch.Tensor, target: torch.Tensor, ratio: int, plan: CSATNDPlan, chunk_size: int,
    ) -> torch.Tensor:
        """Calculate native student gradients and the original scalar KL objective."""
        options, metadata = _tnd_kl_metadata(plan, query.device, query.shape[1], query.shape[2], ratio)
        kernel, _ = _load_kl_ops()
        with torch.autocast(device_type=query.device.type, enabled=False):
            dq, dk, dw, prediction = kernel(
                query, key.unsqueeze(1), weights.float(), indices, target.unsqueeze(1), metadata=metadata, **options,
            )
            logs = []
            q_start = k_start = 0
            for q_length, k_length in zip(plan.q_lengths, plan.k_lengths):
                q_slice = slice(q_start, q_start+q_length)
                logs.append(_stable_prediction_log(
                    prediction[q_slice, 0].unsqueeze(0), query[q_slice].unsqueeze(0),
                    key[k_start:k_start+k_length].unsqueeze(0), weights[q_slice].unsqueeze(0),
                    indices[q_slice, 0].unsqueeze(0), chunk_size,
                )[0])
                q_start += q_length
                k_start += k_length
            log_prediction = torch.cat(logs)
            loss = (target * (target.clamp_min(torch.finfo(torch.float32).tiny).log() - log_prediction)).sum()
        ctx.save_for_backward(dq, dk.squeeze(1), dw)
        return loss

    @staticmethod
    def backward(ctx: Any, seed: torch.Tensor) -> tuple:
        """Apply the upstream scalar seed to all three precomputed student gradients."""
        return (*(value.float() * seed for value in ctx.saved_tensors), None, None, None, None, None)


def sparse_lightning_indexer_kl_loss(
        query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor, attention_query: torch.Tensor,

    Returns:
        Scalar KL normalized by all local queries, including empty-prefix rows.
        The external teacher retains its per-document bounded computation.
    """
    if len(segments) == 1:
        return _bsnd_sparse_lightning_indexer_kl_loss(
            query, key, weights, attention_query, main_key, indices, sinks,
            scale, coefficient, chunk_size, ratio, segments[0],
        )
    if query.shape[0] != 1:
        raise ValueError("Packed Lightning Indexer KL requires batch size one")
    total = query[..., :0].float().sum() + key[..., :0].float().sum() + weights[..., :0].float().sum()
    plan = csa_tnd_plan(segments, ratio)
    if plan is None:
        return total
    selected = plan.local_indices(indices, padded_width=512)
    targets = []
    cursor = 0
    with torch.no_grad(), torch.autocast(device_type=query.device.type, enabled=False):
        for start, end, key_start, key_end, residual in plan.segments:
            targets.append(causal_attention_teacher(
                attention_query[:, :, start:end], main_key[:, key_start:key_end],
                selected[cursor:cursor+end-start, 0].unsqueeze(0), sinks, scale, ratio, residual, chunk_size,
            )[0])
            cursor += end-start
    loss = _SparseLightningIndexerKLLossTND.apply(
        plan.query(query), plan.key(key), plan.query(weights), selected,
        torch.cat(targets).contiguous(), ratio, plan, chunk_size,
    )
    return total + loss * (coefficient / query.shape[1])

        device = torch.device(device)
        # Torch exposes tensor mutation tracking only through its version counter.
        version = None if self.cu_seq_lens.is_inference() else self.cu_seq_lens._version  # pylint: disable=protected-access
        if version is not None and self._snapshot.get("version") == version:
            prepared: SharedCompressedPackedSequence = self if self._snapshot.get("owned") else self._snapshot["value"]
        else:
            prepared = SharedCompressedPackedSequence(
                self.cu_seq_lens.detach().clone(), self.local_query_start,
                self.local_query_length, self.global_sequence_length,

        return self._prepared_positions[device]

    def segment_start_mask(self, device: torch.device) -> torch.Tensor:
        """Reuse the per-batch sample-start mask across model forwards."""
        device = torch.device(device)
        if device not in self._prepared_masks:
            self._prepared_masks[device] = self.local_positions(device) == self.local_segment_starts(device)
        return self._prepared_masks[device]

    def indexer_segments(self, ratio: int) -> tuple[tuple[int, int, int, int, int], ...]:
        """Describe packed Q intervals and completed K prefixes without per-layer synchronization."""
        if ratio not in self._indexer_segments:

            device: Device on which tensors and prepared geometry are consumed.
            compress_ratio: Compression ratio of the consuming layer.
            segment_starts: Optional precomputed global sample starts for local queries.
        """
        if compress_ratio <= 0:
            raise ValueError("minimum_key_indices requires a positive compression ratio")
        cached = self._prepared_minima.get((torch.device(device), compress_ratio))
        if cached is not None:
            return cached
        self.validate_compression_alignment(compress_ratio)
        if segment_starts is None:
            segment_starts = self.local_segment_starts(device)
        return segment_starts // compress_ratio

    def local_segment_starts(self, device: torch.device) -> torch.Tensor:
        """Return each local query's global packed-sample start position.

            device: Device on which tensors and prepared geometry are consumed.
        """
        cached = self._prepared_starts.get(torch.device(device))
        if cached is not None:
            return cached
        if self.local_query_start < 0 or self.local_query_length < 0 or (
                self.local_query_start + self.local_query_length > self.global_sequence_length
        ):
            raise ValueError("packed local query interval must lie within the global sequence")
        boundaries = self.cu_seq_lens.to(dtype=torch.long)
        if boundaries.ndim != 1 or boundaries.numel() < 2 or int(boundaries[0]) != 0:
            raise ValueError("packed cu_seq_lens must be one-dimensional and start with zero")
        if int(boundaries[-1]) != self.global_sequence_length:

        raise ValueError("Fused LightningIndexer requires query_segments when minimum_key_indices is set")
    validate_indexer_tensors(query, key, merge_weight)
    if query.device.type != "npu":
        raise RuntimeError(f"Fused LightningIndexer requires NPU inputs; got device={query.device}")
    if query_segments is None:
        length = query.shape[1]
        raw_end = query_offset + length
        query_segments = ((0, length, 0, raw_end // compress_ratio, raw_end % compress_ratio),)
    return lightning_indexer(query, key, merge_weight, compress_ratio, top_k, query_segments)


@torch.no_grad()
def compressed_causal_topk(

        raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
    if query_chunk_size <= 0:
        raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
    if query.ndim != 4:
        raise ValueError(f"Indexer query must have shape [B, Q, H, D], got {tuple(query.shape)}")
    batch_size, sequence_length, _, head_dim = query.shape
    if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
        raise ValueError(
            "compressed index key must have shape [batch, compressed_sequence, head_dim]"

        raise ValueError(f"compress_ratio must be positive, got {compress_ratio}")
    if query_chunk_size <= 0:
        raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
    batch_size, sequence_length, _, head_dim = query.shape
    if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
        raise ValueError(
            "compressed index key must have shape [batch, compressed_sequence, head_dim]"
        )
    if merge_weight.shape != query.shape[:3]:

        raise ValueError(f"block_size must be positive, got {block_size}")
    if query_chunk_size <= 0:
        raise ValueError(f"query_chunk_size must be positive, got {query_chunk_size}")
    batch_size, sequence_length, _, head_dim = query.shape
    if key.ndim != 3 or key.shape[::2] != (batch_size, head_dim):
        raise ValueError(
            "compressed index key must have shape [batch, compressed_sequence, head_dim]"
        )
    if merge_weight.shape != query.shape[:3]:

    if topk_indices.shape[:2] != index_query.shape[:2]:
        raise ValueError("topk_indices and index_query must share batch/query dimensions")
    if use_fused:
        if tp_context is not None:
            raise NotImplementedError("Fused KL does not support TP score/teacher reduction")
        if query_segments is None:
            raise ValueError("Fused KL requires document-local causal query_segments")
        validate_fused_kl_inputs(index_query, index_key, merge_weight, topk_indices)
        return sparse_lightning_indexer_kl_loss(
            index_query, index_key, merge_weight, attention_query.detach(), compressed_key.detach(),
            topk_indices, sinks.detach(), attention_scale, loss_coeff, query_chunk_size,
            compress_ratio, query_segments,
        )

        merge_weight = merge_weight * (self.head_dim**-0.5 * self.global_num_heads**-0.5)
        reduce_sum = None if tp_context is None else tp_context.reduce_sum
        if key_handle is not None:
            key = key_handle.wait()
        topk_indices, candidate_blocks = self._select_indices(
            query, key, merge_weight, candidate_blocks, query_offset,
            reduce_sum, minimum_key_indices, query_segments,
        )
        return SharedCompressedIndexerOutput(
            topk_indices=topk_indices,
            index_key=key,
            candidate_blocks=candidate_blocks,
            index_query=query,

                minimum_key_indices=minimum_key_indices,
                query_segments=query_segments,
                use_fused=self.use_fused,
            )
        return topk_indices, candidate_blocks


def reference_sparse_attention(
        query: torch.Tensor,

    ) -> tuple[torch.Tensor | None, torch.Tensor | None, tuple]:
        """Validate packed document boundaries and derive local Q/K segments."""
        segment_starts = None
        minimum_key_indices = None
        ratio = max(1, self.compress_ratio)
        raw_end = query_offset + sequence_length
        query_segments = ((0, sequence_length, 0, raw_end // ratio, raw_end % ratio),)
        if packed_sequence is not None:
            if batch_size != 1:
                raise ValueError("V4.1 compact packed attention currently requires micro_batch_size=1")
            if (

                    packed_sequence.local_query_start != query_offset
                    or packed_sequence.local_query_length != sequence_length
                    or packed_sequence.global_sequence_length != global_sequence_length
            ):
                packed_shard = (
                    packed_sequence.local_query_start,
                    packed_sequence.local_query_length,
                    packed_sequence.global_sequence_length,
                )

                    "packed sequence geometry does not match the local CP shard: "
                    f"packed={packed_shard}, "
                    f"attention={(query_offset, sequence_length, global_sequence_length)}"
                )
            packed_sequence = packed_sequence.prepare(device, (self.compress_ratio,))
            segment_starts = packed_sequence.local_segment_starts(device)
            if self.compress_ratio:
                minimum_key_indices = packed_sequence.minimum_key_indices(device, self.compress_ratio)
            query_segments = packed_sequence.indexer_segments(ratio)
        return segment_starts, minimum_key_indices, query_segments

    def _compute_compressed_kv(
            self,
            hidden_states: torch.Tensor,

            window.masked_fill_(window < segment_starts.unsqueeze(-1), -1)
        if compressed_kv is None:
            return key_value, window
        combined_key_value = torch.cat((key_value, compressed_kv.unsqueeze(1)), dim=2)
        compressed_indices = torch.where(topk_indices >= 0, topk_indices.long() + global_sequence_length, -1)
        return combined_key_value, torch.cat((window, compressed_indices), dim=-1)

    def _apply_indexer_loss(
            self,
            indexer_output: Any | None,

                sequence_length,
                self.compress_ratio,
                self.is_kv_source,
            )
        segment_starts, minimum_key_indices, query_segments = self._resolve_packed_segments(
            packed_sequence,
            batch_size,
            sequence_length,
            query_offset,

            shared_state.publish_compressed_kv(self.layer_idx, compressed_handle.wait())
        if key_value is None:
            raise RuntimeError("raw KV all-gather did not produce a tensor")

        compressed_kv = None
        topk_indices = None
        if self.compress_ratio:
            compressed_kv = shared_state.require_compressed_kv(self.kv_source_layer_idx, self.layer_idx)
            topk_indices = shared_state.require_topk_indices(self.index_source_layer_idx, self.layer_idx)
        query = self._apply_indexer_loss(indexer_output, query, compressed_kv, tp_context, query_segments)
        attention_output = self._run_sparse_attention(
            query, key_value, compressed_kv, topk_indices, query_offset, global_sequence_length,
            segment_starts, query_segments,
        )