Diff Coverage

        Flattened mixed residual streams.
    """
    x_shape = x.size()
    residual = residual.reshape(x_shape[0], x_shape[1], num_stream, -1)
    if x.device.type == "npu" and omni_training_custom_ops is not None:
        return _MhcPost.apply(residual, h_res, x, h_post).flatten(2)
    output = h_post.unsqueeze(-1) * x.unsqueeze(-2)
    output = output + torch.sum(
        h_res.unsqueeze(-1) * residual.unsqueeze(-2), dim=2
    )
    return output.flatten(2).to(x.dtype)

        Args:
            assets: Prepared tokenizer map and per-layer hash layout.
            layer_id: Decoder layer owning this Engram table.
        """
        super().__init__()
        layer_ids = tuple(int(value) for value in assets["layer_ids"])
        if layer_id not in layer_ids:
            raise ValueError(f"Engram layer {layer_id} is absent from assets layer_ids={layer_ids}")
        layer_index = layer_ids.index(layer_id)
        self.layer_id = layer_id
        self.max_ngram_size = int(assets["max_ngram_size"])
        self.num_heads = int(assets["num_heads"])
        token_map = torch.tensor(assets["token_map"], dtype=torch.long)
        self.pad_id = int(token_map[int(assets["pad_token_id"])])
        primes = torch.tensor(assets["primes"][layer_index], dtype=torch.long)
        if tuple(primes.shape) != (self.max_ngram_size - 1, self.num_heads):
            raise ValueError(
                "Engram prime layout must be [max_ngram_size - 1, num_heads], "
                f"got {tuple(primes.shape)}"
            )
        flattened_primes = primes.flatten()
        offsets = torch.cat(
            [torch.zeros(1, dtype=torch.long), flattened_primes.cumsum(0)[:-1]]
        ).view_as(primes)
        multipliers = torch.tensor(assets["multipliers"][layer_index], dtype=torch.long)
        if multipliers.numel() != self.max_ngram_size:
            raise ValueError("Engram multiplier count must equal max_ngram_size")
        self.logical_num_embeddings = int(flattened_primes.sum().item())
        initial_buffers = {
            "token_map": token_map,
            "primes": primes,
            "offsets": offsets,
            "multipliers": multipliers,

            "primes": primes,
            "offsets": offsets,
            "multipliers": multipliers,
        }
        for name, value in initial_buffers.items():
            register_rebuildable_buffer(self, name, value=value, persistent=False)

    def forward(
        self,
        input_ids: torch.Tensor,

        segment_starts: torch.Tensor | None = None,
        token_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Return global flat-table rows shaped ``[batch, seq, hash_columns]``."""
        if input_ids.ndim != 2:
            raise ValueError(f"Engram expects input_ids [B,S], got {tuple(input_ids.shape)}")
        compressed = self.token_map[input_ids]
        if token_mask is not None:
            if token_mask.shape != input_ids.shape:
                raise ValueError("Engram token_mask must have the same shape as input_ids")
            compressed = torch.where(
                token_mask.to(torch.bool),
                compressed,
                compressed.new_full((), self.DEAD_TOKEN),
            )
        batch_size, sequence_length = compressed.shape
        positions = torch.arange(sequence_length, device=input_ids.device).expand(batch_size, -1)
        segment_ids = None
        if segment_starts is not None:
            if segment_starts.shape != input_ids.shape:
                raise ValueError("Engram segment_starts must have the same shape as input_ids")
            starts = segment_starts.to(torch.bool).clone()
            starts[:, 0] = True
            segment_ids = starts.to(torch.long).cumsum(dim=1)

        tokens = []
        blocked = torch.zeros_like(positions, dtype=torch.bool)
        for shift in range(self.max_ngram_size):
            source_positions = (positions - shift).clamp_min(0)
            source = compressed.gather(1, source_positions)
            blocked = blocked | (positions < shift)
            if segment_ids is not None:
                source_segments = segment_ids.gather(1, source_positions)
                blocked = blocked | (source_segments != segment_ids)
            blocked = blocked | (source == self.DEAD_TOKEN)
            tokens.append(torch.where(blocked, self.pad_id, source))

        products = torch.stack(tokens, dim=-1) * self.multipliers
        rolling = products[..., 0]
        hashes = []
        for ngram_index in range(self.max_ngram_size - 1):
            rolling = torch.bitwise_xor(rolling, products[..., ngram_index + 1])
            row = rolling.unsqueeze(-1) % self.primes[ngram_index]
            hashes.append(row + self.offsets[ngram_index])
        return torch.cat(hashes, dim=-1)


def _all_gather_sequence(value: torch.Tensor, group: Any, group_size: int) -> torch.Tensor:
    """Gather equal-size sequence chunks and concatenate along dimension one."""
    gathered = [torch.empty_like(value) for _ in range(group_size)]
    dist.all_gather(gathered, value.contiguous(), group=group)
    return torch.cat(gathered, dim=1)


@module_replacement
class EngramModule(nn.Module):

        module_fqn: str = "",
        context: Mapping[str, Any] | None = None,
    ) -> None:
        """Reuse a source Engram module's parameters and hash mapping."""
        super().__init__()
        del module_fqn, context
        required = ("hash_mapping", "embed", "wkv", "q_weight", "k_weight")
        missing = [name for name in required if not hasattr(module, name)]
        if missing:
            raise TypeError(f"Engram source is missing required attributes: {missing}")
        self.hash_mapping = module.hash_mapping
        self.embed = module.embed
        self.wkv = module.wkv
        self.q_weight = module.q_weight
        self.k_weight = module.k_weight
        self.layer_id = int(module.layer_id)
        self.hidden_size = int(module.hidden_size)
        self.hc_mult = int(module.hc_mult)
        self.eps = float(module.eps)
        self.clamp_value = float(module.clamp_value)
        self.logical_num_embeddings = int(module.logical_num_embeddings)
        self.padded_num_embeddings = int(module.padded_num_embeddings)
        self.train(module.training)

    def _fuse(self, hidden_states: torch.Tensor, embeddings: torch.Tensor) -> torch.Tensor:
        """Apply the exact V4.1 WKV projection and signed-root gate."""
        key_value = self.wkv(embeddings.flatten(start_dim=-2))
        key, value = key_value.split(
            [self.hc_mult * self.hidden_size, self.hidden_size], dim=-1
        )
        key = key.float().unflatten(-1, (self.hc_mult, self.hidden_size))
        hidden_fp32 = hidden_states.float()
        weight = self.q_weight.float() * self.k_weight.float()
        reciprocal_std = torch.rsqrt(hidden_fp32.square().mean(-1) + self.eps)
        reciprocal_std = reciprocal_std * torch.rsqrt(key.square().mean(-1) + self.eps)
        dot = (hidden_fp32 * weight * key).sum(-1)
        dot = dot * reciprocal_std * self.hidden_size**-0.5
        root = dot.abs().clamp_min(self.clamp_value).sqrt()
        signed_root = torch.where(dot < 0, -root, root)
        gate = torch.sigmoid(signed_root)
        return (
            hidden_fp32 + gate.unsqueeze(-1) * value.float().unsqueeze(-2)
        ).to(hidden_states.dtype)

    @staticmethod

        ep_size: int,
        padded_num_embeddings: int,
    ) -> torch.Tensor:
        """Route requested global rows to their EP owners and restore order."""
        if padded_num_embeddings % ep_size:
            raise ValueError(
                f"Engram padded rows {padded_num_embeddings} are not divisible by ep_size={ep_size}"
            )
        num_local_rows = padded_num_embeddings // ep_size
        row_start = ep_rank * num_local_rows
        flat_ids = hash_ids.reshape(-1)
        owner = torch.div(flat_ids, num_local_rows, rounding_mode="floor")
        if owner.numel() and int(owner.max()) >= ep_size:
            raise IndexError("Engram hash requested a row outside the padded table")
        local_counts = torch.bincount(owner, minlength=ep_size).to(torch.long)
        gathered_counts = [torch.empty_like(local_counts) for _ in range(ep_size)]
        dist.all_gather(gathered_counts, local_counts, group=ep_group)
        all_counts = torch.stack(gathered_counts)
        send_counts = local_counts.cpu().tolist()
        recv_counts = all_counts[:, ep_rank].cpu().tolist()

        # Ascend ArgSort executes integer inputs on AiCPU. The owner range is
        # bounded by ep_size, so float32 preserves it exactly and stays on AiCore.
        sort_indices = torch.argsort(owner.float(), stable=True)
        sorted_ids = flat_ids.index_select(0, sort_indices)
        received_ids = ep_all_to_all(sorted_ids, send_counts, recv_counts, ep_group)
        local_ids = received_ids - row_start
        if local_ids.numel() and (int(local_ids.min()) < 0 or int(local_ids.max()) >= num_local_rows):
            raise IndexError("Engram EP routing delivered a row to the wrong owner")
        # Call the child module instead of reading ``embedding.weight``
        # directly. This preserves the FSDP pre/post-forward lifecycle when
        # the EP table is wrapped independently from the dense Engram gates.
        values = embedding(local_ids)
        returned = ep_all_to_all(values, recv_counts, send_counts, ep_group)
        restored = torch.zeros_like(returned)
        restored.scatter_add_(
            0,
            sort_indices.unsqueeze(1).expand(-1, returned.shape[1]),
            returned,
        )
        return restored.view(*hash_ids.shape, embedding.embedding_dim)

    def _aligned_hash_ids(
        self,
        hidden_states: torch.Tensor,

        tp_rank: int = 0,
        tp_size: int = 1,
    ) -> torch.Tensor:
        """Hash the local sequence plus the left halo required by n-grams."""
        cp_local_ids = input_ids
        if cp_size > 1:
            input_ids = _all_gather_sequence(input_ids, cp_group, cp_size)
            if segment_starts is not None:
                segment_starts = _all_gather_sequence(segment_starts, cp_group, cp_size)
            if token_mask is not None:
                token_mask = _all_gather_sequence(token_mask, cp_group, cp_size)
        local_sequence = hidden_states.shape[1]
        cp_sequence = cp_local_ids.shape[1]
        if cp_sequence == local_sequence:
            sequence_parallel_parts = 1
            local_tp_rank = 0
        elif cp_sequence % local_sequence == 0 and cp_sequence // local_sequence == tp_size:
            sequence_parallel_parts = tp_size
            local_tp_rank = tp_rank
        else:
            raise ValueError(
                "Engram input_ids length is incompatible with the local hidden sequence: "
                f"ids={cp_sequence}, hidden={local_sequence}, tp_size={tp_size}"
            )
        del sequence_parallel_parts
        start = cp_rank * cp_sequence + local_tp_rank * local_sequence
        end = start + local_sequence
        context_start = max(0, start - self.hash_mapping.max_ngram_size + 1)
        ids_window = input_ids[:, context_start:end]
        starts_window = None if segment_starts is None else segment_starts[:, context_start:end]
        mask_window = None if token_mask is None else token_mask[:, context_start:end]
        hashes = self.hash_mapping(ids_window, starts_window, mask_window)
        target_offset = start - context_start
        return hashes[:, target_offset:target_offset + local_sequence]

    def forward(
        self,
        hidden_states: torch.Tensor,

        segment_starts: torch.Tensor | None = None,
        token_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Run replicated single-card/EP1 Engram."""
        hash_ids = self._aligned_hash_ids(hidden_states, input_ids, segment_starts, token_mask)
        fused = self._fuse(hidden_states, self.embed(hash_ids))
        if token_mask is None:
            return fused
        return torch.where(token_mask.to(torch.bool).unsqueeze(-1).unsqueeze(-1), fused, hidden_states)

    def parallel_forward(
        self,
        hidden_states: torch.Tensor,

        tp_rank: int = 0,
        tp_size: int = 1,
    ) -> torch.Tensor:
        """Run CP/TP sequence alignment and sparse EP lookup."""
        hash_ids = self._aligned_hash_ids(
            hidden_states,
            input_ids,
            segment_starts,
            token_mask,

            cp_size=cp_size,
            tp_rank=tp_rank,
            tp_size=tp_size,
        )
        embeddings = self._sparse_lookup(
            hash_ids,
            self.embed,
            ep_group=ep_group,
            ep_rank=ep_rank,

            ep_rank=ep_rank,
            ep_size=ep_size,
            padded_num_embeddings=self.padded_num_embeddings,
        )
        fused = self._fuse(hidden_states, embeddings)
        if token_mask is None:
            return fused
        return torch.where(token_mask.to(torch.bool).unsqueeze(-1).unsqueeze(-1), fused, hidden_states)


__all__ = ["EngramModule", "NgramHashMapping"]

        module_fqn: str = "",
        context: Mapping[str, Any] | None = None,
    ) -> None:
        """Reuse the source module's ``fn/base/scale`` parameter layout."""
        super().__init__()
        del module_fqn, context
        required = ("input_norm", "fn", "base", "scale", "hc_mult")
        missing = [name for name in required if not hasattr(module, name)]
        if missing:
            raise TypeError(f"pipelined MHC source is missing required attributes: {missing}")
        self.input_norm = module.input_norm
        self.fn = module.fn
        self.base = module.base
        self.scale = module.scale
        self.hc_mult = int(module.hc_mult)
        self.hc_sinkhorn_iters = int(module.hc_sinkhorn_iters)
        self.hc_eps = float(module.hc_eps)
        expected_mix = (self.hc_mult + 2) * self.hc_mult
        if tuple(self.fn.shape)[0] != expected_mix or self.base.numel() != expected_mix:
            raise ValueError("pipelined MHC fn/base shapes do not match hc_mult")
        if self.scale.numel() != 3:
            raise ValueError("pipelined MHC scale must contain three values")
        self.train(module.training)

    def forward(
        self,
        hidden_streams: torch.Tensor,

        self,
        hidden_streams: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Produce the coefficients consumed by the following sublayer."""
        flattened = self.input_norm(hidden_streams.flatten(start_dim=2).float())
        mix = F.linear(  # pylint: disable=not-callable
            flattened, self.fn.float()
        )
        num_stream = self.hc_mult
        pre, post, residual = mix.split(
            [num_stream, num_stream, num_stream * num_stream], dim=-1
        )
        pre_bias, post_bias, residual_bias = self.base.split(
            [num_stream, num_stream, num_stream * num_stream]
        )
        pre_scale, post_scale, residual_scale = self.scale.unbind(0)
        pre = torch.sigmoid(pre * pre_scale + pre_bias) + self.hc_eps
        post = 2 * torch.sigmoid(post * post_scale + post_bias)
        residual = residual.view(*residual.shape[:-1], num_stream, num_stream)
        residual = residual * residual_scale + residual_bias.view(num_stream, num_stream)
        if hidden_streams.device.type == "npu" and hasattr(torch.ops.custom, "npu_sinkhorn"):
            residual = sinkhorn(residual, self.hc_sinkhorn_iters, self.hc_eps)
        else:
            residual = sinkhorn_knopps(residual, self.hc_sinkhorn_iters, self.hc_eps)
        return pre, post, residual


def pipelined_mhc_post(
    sublayer_output: torch.Tensor,

    post: torch.Tensor,
    combine: torch.Tensor,
) -> torch.Tensor:
    """Apply the shared high-performance mHC post path to 4-D streams."""
    num_stream = residual.shape[-2]
    flattened = mhc_post(
        sublayer_output,
        residual.flatten(start_dim=2),
        post,
        combine,

        post,
        combine,
        num_stream,
    )
    return flattened.unflatten(-1, (num_stream, residual.shape[-1]))

    gather_sequence: Callable[[torch.Tensor, int], torch.Tensor]

    def wait(self) -> torch.Tensor:
        """Run the fallback gather at the first consumer boundary."""
        return self.gather_sequence(self.tensor, self.sequence_dim)


@dataclass
class SharedCompressedAttentionState:

            value_name: str,
            consumer_layer: int,
    ) -> torch.Tensor:
        """Return one published tensor or report the broken layer dependency."""
        if source_layer is None or source_layer not in values:
            raise RuntimeError(
                f"layer {consumer_layer} requires {value_name} from source "
                f"layer {source_layer}, but that source has not run in this forward"
            )
        return values[source_layer]

    def publish_compressed_kv(self, source_layer: int, value: torch.Tensor) -> None:
        """Publish compressed K=V while preserving its autograd graph."""
        self.compressed_kv_by_source[source_layer] = value

    def require_compressed_kv(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
        """Read compressed K=V from the consumer's configured Full layer."""
        return self._require(
            self.compressed_kv_by_source,
            source_layer,
            "compressed K=V",
            consumer_layer,

        )

    def publish_index_key(self, source_layer: int, value: torch.Tensor) -> None:
        """Publish the Full Indexer's shared key tensor."""
        self.index_key_by_source[source_layer] = value

    def require_index_key(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
        """Read the Indexer key associated with a compressed-KV source."""
        return self._require(
            self.index_key_by_source,
            source_layer,
            "Indexer key",
            consumer_layer,

        )

    def publish_topk_indices(self, source_layer: int, value: torch.Tensor) -> None:
        """Publish one Full/Reindex layer's token selection."""
        self.topk_indices_by_source[source_layer] = value

    def require_topk_indices(self, source_layer: int | None, consumer_layer: int) -> torch.Tensor:
        """Read the latest configured Full/Reindex selection."""
        return self._require(
            self.topk_indices_by_source,
            source_layer,
            "Top-K indices",
            consumer_layer,

        )

    def publish_candidate_blocks(self, source_layer: int, value: torch.Tensor) -> None:
        """Publish hierarchical candidate blocks from the configured source."""
        self.candidate_blocks_by_source[source_layer] = value

    def require_candidate_blocks(
            self,
            source_layer: int | None,

            source_layer: int | None,
            consumer_layer: int,
    ) -> torch.Tensor:
        """Read hierarchical candidates for a later Reindex layer."""
        return self._require(
            self.candidate_blocks_by_source,
            source_layer,
            "candidate blocks",
            consumer_layer,

    launch_sequence: Callable[[torch.Tensor, int], SequenceGatherHandle] | None = None

    def launch(self, tensor: torch.Tensor, sequence_dim: int) -> SequenceGatherHandle:
        """Launch an async gather or create a deferred synchronous gather."""
        if self.launch_sequence is not None:
            return self.launch_sequence(tensor, sequence_dim)
        return _DeferredSequenceGather(tensor, sequence_dim, self.gather_sequence)


@dataclass(frozen=True)
class SharedCompressedAttentionTPContext:

    global_sequence_length: int

    def local_segment_starts(self, device: torch.device) -> torch.Tensor:
        """Return each local query's global packed-sample start position."""
        boundaries = self.cu_seq_lens.to(device=device, 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(
                "packed cu_seq_lens must cover global_sequence_length: "
                f"{int(boundaries[-1])} versus {self.global_sequence_length}"
            )
        if torch.any(boundaries[1:] <= boundaries[:-1]):
            raise ValueError("packed cu_seq_lens must be strictly increasing")
        query_positions = torch.arange(
            self.local_query_start,
            self.local_query_start + self.local_query_length,
            device=device,
        )
        segment_ids = torch.bucketize(query_positions, boundaries[1:], right=True)
        return boundaries[:-1].index_select(0, segment_ids).unsqueeze(0)

    def validate_compression_alignment(self, compress_ratio: int) -> None:
        """Reject packed samples whose compressor groups would cross boundaries."""
        if compress_ratio <= 1:
            return
        boundaries = self.cu_seq_lens.to(dtype=torch.long)
        misaligned = boundaries[boundaries.remainder(compress_ratio) != 0]
        if misaligned.numel():
            raise ValueError(
                "packed sample boundaries must align with the CSA2 compression ratio "
                f"{compress_ratio}; misaligned boundaries={misaligned.tolist()}"
            )

def _rotate_half(tensor: torch.Tensor) -> torch.Tensor:
    """Rotate adjacent real/imaginary pairs for V4 interleaved RoPE."""
    first = tensor[..., 0::2]
    second = tensor[..., 1::2]
    return torch.stack((-second, first), dim=-1).flatten(-2)


def _apply_v41_rope(
        tensor: torch.Tensor,

        sin: torch.Tensor,
        unsqueeze_dim: int = 1,
) -> torch.Tensor:
    """Apply V4 interleaved RoPE to the trailing rotary channels."""
    cos = cos.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
    sin = sin.repeat_interleave(2, dim=-1).unsqueeze(unsqueeze_dim)
    rope_dim = cos.shape[-1]
    pass_through, rotary = tensor[..., :-rope_dim], tensor[..., -rope_dim:]
    rotary = ((rotary.float() * cos) + (_rotate_half(rotary).float() * sin)).to(tensor.dtype)
    return torch.cat((pass_through, rotary), dim=-1)


def build_sliding_window_indices(
        batch_size: int,

        query_offset: int = 0,
        key_length: int | None = None,
) -> torch.Tensor:
    """Build fixed-width causal sliding-window indices."""
    if key_length is None:
        key_length = sequence_length
    width = min(key_length, window_size)
    query = torch.arange(
        query_offset,
        query_offset + sequence_length,
        device=device,
    ).unsqueeze(-1)
    offsets = torch.arange(width - 1, -1, -1, device=device)
    indices = query - offsets
    indices = indices.masked_fill(indices < 0, -1)
    return indices.unsqueeze(0).expand(batch_size, -1, -1)


def compressed_causal_topk(
        query: torch.Tensor,

    groups of ``compress_ratio`` source tokens. This implementation retains
    the source rule ``key < floor((query + 1) / ratio)`` while using batched
    matrix multiplication and bounded query chunks on accelerator cores.
    """
    if compress_ratio <= 0:
        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, num_heads, head_dim = query.shape
    if key.ndim != 3 or key.shape[0] != batch_size or key.shape[2] != 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"merge_weight must have shape {tuple(query.shape[:3])}, got {tuple(merge_weight.shape)}"
        )
    compressed_length = key.shape[1]
    if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
        raise ValueError("minimum_key_indices must have shape [batch, query]")
    top_k = min(sparse_count, compressed_length)
    if top_k <= 0:
        return torch.empty(
            batch_size,
            sequence_length,
            0,
            dtype=torch.int32,

            dtype=torch.int32,
            device=query.device,
        )

    key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
    key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
    selected_chunks = []
    for start in range(0, sequence_length, query_chunk_size):
        end = min(start + query_chunk_size, sequence_length)
        scores = torch.matmul(query[:, start:end].float(), key_fp32).relu_()
        scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
        if reduce_sum is not None:
            scores = reduce_sum(scores)
        visible = (
            torch.arange(query_offset + start, query_offset + end, device=query.device) + 1
        ) // compress_ratio
        scores.masked_fill_(key_positions >= visible.view(1, -1, 1), float("-inf"))
        if minimum_key_indices is not None:
            minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
            scores.masked_fill_(key_positions < minimum, float("-inf"))
        top = scores.topk(top_k, dim=-1, sorted=False)
        indices = top.indices.masked_fill(~torch.isfinite(top.values), compressed_length)
        indices = indices.sort(dim=-1).values
        indices = indices.masked_fill(indices == compressed_length, -1)
        selected_chunks.append(indices.to(torch.int32))
    return torch.cat(selected_chunks, dim=1)


def select_candidate_blocks(
        logits: torch.Tensor,

    The newest reachable partial block is pinned into the candidate set, then
    the remaining blocks compete by their maximum position score. ``logits``
    must already contain ``-inf`` at causally unreachable positions.
    """
    if topk_blocks <= 0:
        raise ValueError(f"topk_blocks must be positive, got {topk_blocks}")
    if block_size <= 0:
        raise ValueError(f"block_size must be positive, got {block_size}")
    width = logits.shape[-1]
    if width == 0:
        return torch.zeros_like(logits, dtype=torch.bool)
    padding = -width % block_size
    scores = torch.nn.functional.pad(logits, (0, padding), value=float("-inf"))
    scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
    num_blocks = scores.shape[-1]
    if isinstance(compress_lens, torch.Tensor):
        last = (compress_lens.to(device=logits.device) - 1) // block_size
    else:
        last = (compress_lens - 1) // block_size
    block_positions = torch.arange(num_blocks, device=logits.device)
    scores = scores.masked_fill(block_positions == last, float("inf"))
    top = scores.topk(min(topk_blocks, num_blocks), dim=-1)
    keep = torch.zeros_like(scores, dtype=torch.bool)
    keep.scatter_(-1, top.indices, top.values > float("-inf"))
    return keep.repeat_interleave(block_size, dim=-1)[..., :width]


def select_candidate_block_indices(
        logits: torch.Tensor,

    Storing block ids uses ``block_size`` times less cross-layer state than
    expanded position ids and avoids the full ``[batch, query, key]`` boolean
    mask used by the released inference reference. Invalid slots are ``-1``.
    """
    if topk_blocks <= 0:
        raise ValueError(f"topk_blocks must be positive, got {topk_blocks}")
    if block_size <= 0:
        raise ValueError(f"block_size must be positive, got {block_size}")
    width = logits.shape[-1]
    if width == 0:
        return torch.empty(*logits.shape[:-1], 0, dtype=torch.int32, device=logits.device)
    padding = -width % block_size
    scores = torch.nn.functional.pad(logits, (0, padding), value=float("-inf"))
    scores = scores.unflatten(-1, (-1, block_size)).amax(dim=-1)
    num_blocks = scores.shape[-1]
    if isinstance(compress_lens, torch.Tensor):
        last = (compress_lens.to(device=logits.device) - 1) // block_size
    else:
        last = (compress_lens - 1) // block_size
    block_positions = torch.arange(num_blocks, device=logits.device)
    scores = scores.masked_fill(block_positions == last, float("inf"))
    top = scores.topk(min(topk_blocks, num_blocks), dim=-1, sorted=False)
    return top.indices.masked_fill(top.values == float("-inf"), -1).to(torch.int32)


def compressed_causal_candidates(
        query: torch.Tensor,

        reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
        minimum_key_indices: torch.Tensor | None = None,
) -> torch.Tensor:
    """Build compact blockwise candidate ids in bounded query chunks."""
    _, sequence_length, _, _ = query.shape
    compressed_length = key.shape[1]
    if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
        raise ValueError("minimum_key_indices must have shape [batch, query]")
    key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
    key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
    candidate_chunks = []
    for start in range(0, sequence_length, query_chunk_size):
        end = min(start + query_chunk_size, sequence_length)
        scores = torch.matmul(query[:, start:end].float(), key_fp32).relu_()
        scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
        if reduce_sum is not None:
            scores = reduce_sum(scores)
        visible = (
            torch.arange(query_offset + start, query_offset + end, device=query.device) + 1
        ) // compress_ratio
        scores.masked_fill_(key_positions >= visible.view(1, -1, 1), float("-inf"))
        if minimum_key_indices is not None:
            minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
            scores.masked_fill_(key_positions < minimum, float("-inf"))
        candidate_chunks.append(
            select_candidate_block_indices(
                scores,
                visible.view(1, -1, 1),
                topk_blocks,

                topk_blocks,
                block_size,
            )
        )
    return torch.cat(candidate_chunks, dim=1)


def compressed_causal_topk_and_candidates(
        query: torch.Tensor,

        reduce_sum: Callable[[torch.Tensor], torch.Tensor] | None = None,
        minimum_key_indices: torch.Tensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Select Full-layer Top-K and candidate blocks from one score pass."""
    if compress_ratio <= 0:
        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[0] != batch_size or key.shape[2] != 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"merge_weight must have shape {tuple(query.shape[:3])}, got {tuple(merge_weight.shape)}"
        )
    compressed_length = key.shape[1]
    if minimum_key_indices is not None and minimum_key_indices.shape != query.shape[:2]:
        raise ValueError("minimum_key_indices must have shape [batch, query]")
    top_k = min(sparse_count, compressed_length)
    if top_k <= 0:
        empty_topk = torch.empty(
            batch_size,
            sequence_length,
            0,
            dtype=torch.int32,

            0,
            dtype=torch.int32,
            device=query.device,
        )
        return empty_topk, empty_topk

    key_fp32 = key.float().transpose(1, 2).unsqueeze(1)
    key_positions = torch.arange(compressed_length, device=query.device).view(1, 1, -1)
    selected_chunks = []
    candidate_chunks = []
    for start in range(0, sequence_length, query_chunk_size):
        end = min(start + query_chunk_size, sequence_length)
        scores = torch.matmul(query[:, start:end].float(), key_fp32).relu_()
        scores = (scores * merge_weight[:, start:end].float().unsqueeze(-1)).sum(dim=2)
        if reduce_sum is not None:
            scores = reduce_sum(scores)
        visible = (
            torch.arange(query_offset + start, query_offset + end, device=query.device) + 1
        ) // compress_ratio
        scores.masked_fill_(key_positions >= visible.view(1, -1, 1), float("-inf"))
        if minimum_key_indices is not None:
            minimum = minimum_key_indices[:, start:end].unsqueeze(-1)
            scores.masked_fill_(key_positions < minimum, float("-inf"))
        top = scores.topk(top_k, dim=-1, sorted=False)
        indices = top.indices.masked_fill(~torch.isfinite(top.values), compressed_length)
        indices = indices.sort(dim=-1).values
        selected_chunks.append(indices.masked_fill(indices == compressed_length, -1).to(torch.int32))
        candidate_chunks.append(
            select_candidate_block_indices(
                scores,
                visible.view(1, -1, 1),
                topk_blocks,