Diff Coverage

    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(

        raise RuntimeError(
            "use_fused_ops=True requires cann_ops_transformer.lightning_indexer_metadata on this device. "
            "Install the ops-transformer PyTorch extension built for the runtime torch-npu and CANN versions."
        ) from err
    return lightning_indexer_metadata


@lru_cache(maxsize=1)
def _load_kl_ops() -> tuple[Callable[..., Any], Callable[..., Any]]:

            "use_fused_ops=True requires cann_ops_transformer.sparse_lightning_indexer_kl_loss_grad "
            "and sparse_lightning_indexer_kl_loss_grad_metadata. Install the ops-transformer PyTorch extension "
            "built for the runtime torch-npu and CANN versions."
        ) from err
    return sparse_lightning_indexer_kl_loss_grad, sparse_lightning_indexer_kl_loss_grad_metadata


def _supported_inputs(query: torch.Tensor, key: torch.Tensor) -> bool:
    """Check device, dtype and dimensions for a native Indexer/KL call."""

    Other devices keep metadata. Dynamic seqused_q/k would require metadata
    on 910B too; neither adapter below supplies those inputs.
    """
    return "910B" not in torch.npu.get_device_name(device)


def supports_fused_indexer_inputs(query: torch.Tensor, key: torch.Tensor, topk: int) -> bool:
    """Check the LI V2 input contract independently of extension loading.

        topk: Requested sparse selection width.
    """
    if not _supported_inputs(query, key) or not 1 <= topk <= 8192 or not 1 <= query.shape[2] <= 64:
        return False
    if topk > 2048 and topk % 1024:
        return False
    return True


@torch.no_grad()
def fused_compressed_topk(

    batch, length, heads, dim = query.shape
    result = torch.full((batch, length, topk), -1, device=query.device, dtype=torch.int32)
    for start, end, key_start, key_end, residual in segments:
        if key_end <= key_start or start == end:
            continue
        residue = torch.full((batch,), residual, device=query.device, dtype=torch.int32)
        metadata = None
        if _requires_metadata(query.device):
            metadata = _load_indexer_metadata_op()(
                heads, 1, dim, topk, cmp_residual_k=residue,
                batch_size=batch, max_seqlen_q=end-start, max_seqlen_k=key_end-key_start,
                layout_q="BSND", layout_k="BSND", mask_mode=3, cmp_ratio=ratio,
            )

    if not _supported_inputs(query, key) or query.shape[2] not in (8, 16, 32, 64):
        return False
    # Native length validation avoids rejecting valid long banks based only
    # on the conservative sequence range in older operator documentation.
    if indices.shape[-1] != min(512, key.shape[1]) or query.shape[0] > 256:
        return False
    return True


@lru_cache(maxsize=128)
def _kl_metadata(

        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 _stable_prediction_log(
        prediction: torch.Tensor, query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor,

    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 fused_compressed_kl_loss(
        query: torch.Tensor, key: torch.Tensor, weights: torch.Tensor,

        ratio: Compression ratio.
        segments: Contiguous document/local-Q intersections and K prefixes.
    """
    # Zero-length views attach zero gradients even when no segment has a key.
    total = query[..., :0].float().sum() + key[..., :0].float().sum() + weights[..., :0].float().sum()
    for start, end, key_start, key_end, residual in segments:
        if key_end <= key_start:
            continue
        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 = _CausalCompressedKLLoss.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,
        )
        total = total + loss * ((end-start) / query.shape[1])
    return total

            "use_fused_ops=True requires cann_ops_transformer.sparse_flash_mla, sparse_flash_mla_grad "
            "and sparse_flash_mla_metadata. Install the ops-transformer PyTorch extension built for the "
            "runtime torch-npu and CANN versions."
        ) from err
    return sparse_flash_mla, 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")
    outputs = []
    for start, end, key_start, key_end, residual in segments:
        document_start = key_start * max(ratio, 1)
        raw_left = max(document_start, query_offset + start - window + 1)
        raw_right = query_offset + end
        raw = raw_key[:, :, raw_left:raw_right].transpose(1, 2).contiguous()
        main = None
        selected = None
        if ratio and key_end > key_start:
            main = main_key[:, key_start:key_end].unsqueeze(2).contiguous()
            original = indices[:, start:end]
            selected = torch.where(original >= 0, original-key_start, -1).int()
            if selected.shape[-1] < 512:
                selected = torch.nn.functional.pad(selected, (0, 512-selected.shape[-1]), value=-1)
            selected = selected.unsqueeze(2).contiguous()
        part = _SparseMla.apply(
            query[:, :, start:end].transpose(1, 2).contiguous(), raw, main, selected, sinks.float(),
            scale, max(ratio, 1) if main is not None else 1, residual if main is not None else 0, window,
        )
        outputs.append(part)
    return torch.cat(outputs, dim=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:

    if not (compress_ratio <= 128 and reduce_sum is None
            and (minimum_key_indices is None or query_segments is not None)
            and supports_fused_indexer_inputs(query, key, top_k)):
        return None
    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 fused_compressed_topk(query, key, merge_weight, compress_ratio, top_k, query_segments)


@torch.no_grad()
def compressed_causal_topk(

                f"device={query.device}, dtype={query.dtype}, Q={tuple(query.shape)}, K={tuple(key.shape)}, "
                f"topk={top_k}, ratio={compress_ratio}, TP={reduce_sum is not None}. "
                "Check the LI V2 extension, input shapes and ordinary causal geometry; no reference fallback was run."
            )
        return selected

    key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
    key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
    selected_chunks = []

        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]:

                f"device={index_query.device}, dtype={index_query.dtype}, Q={tuple(index_query.shape)}, "
                f"K={tuple(index_key.shape)}, topk={topk_indices.shape[-1]}, TP={tp_context is not None}. "
                "Check the KL extension, head dimensions, TopK512 and causal segments; no reference fallback was run."
            )
        return fused_compressed_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,
        )