perf(components): gather compressed keys with index_select, not two index tensors

hyper_parallel/components/modules/shared_compressed_dsa_attention.py

@@ -274,6 +274,30 @@ def build_sliding_window_indices(
     return indices.unsqueeze(0).expand(batch_size, -1, -1)
 
 
+def _gather_keys(key: torch.Tensor, batch_indices: torch.Tensor,
+                 positions: torch.Tensor) -> torch.Tensor:
+    """Gather compressed keys at ``positions`` for each batch entry.
+
+    Advanced indexing with two index tensors has no AI Core implementation and
+    falls back to AI CPU, where it dominates both the candidate chain and the
+    Indexer KL loss. A single ``index_select`` over the flattened key does the
+    same gather on device, so the batch offset is folded into the index instead
+    of being a second index tensor.
+
+    Args:
+        key: Compressed keys, ``[batch, compressed_length, head_dim]``.
+        batch_indices: Batch ids broadcastable against ``positions``.
+        positions: Key positions to gather, already clamped in range.
+
+    Returns:
+        The gathered keys, ``positions.shape + (head_dim,)``.
+    """
+    compressed_length, head_dim = key.shape[1], key.shape[2]
+    flat_index = (batch_indices * compressed_length + positions).reshape(-1)
+    gathered = key.reshape(-1, head_dim).index_select(0, flat_index)
+    return gathered.view(*positions.shape, head_dim)
+
+
 def compressed_causal_topk(
         query: torch.Tensor,
         key: torch.Tensor,
@@ -298,7 +322,7 @@ 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}")
-    batch_size, sequence_length, num_heads, head_dim = query.shape
+    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]"
@@ -597,7 +621,7 @@ def compressed_candidate_topk(
         if minimum_key_indices is not None:
             valid = valid & (positions >= minimum_key_indices[:, start:end].unsqueeze(-1))
         safe_positions = positions.clamp(min=0, max=max(compressed_length - 1, 0))
-        selected_key = key[batch_indices, safe_positions].float()
+        selected_key = _gather_keys(key, batch_indices, safe_positions).float()
         dots = torch.einsum(
             "bchd,bckd->bchk",
             query[:, start:end].float(),
@@ -655,7 +679,8 @@ def forward(
 
                 query_chunk = index_query[:, start:end].float()
                 weight_chunk = merge_weight[:, start:end].float()
-                selected_index_key = index_key[batch_indices, safe_indices].float()
+                selected_index_key = _gather_keys(
+                    index_key, batch_indices, safe_indices).float()
                 index_dots = torch.einsum(
                     "bcid,bckd->bcik", query_chunk, selected_index_key
                 )
@@ -665,7 +690,8 @@ def forward(
                     index_scores = reduce_sum(index_scores)
                 index_scores.masked_fill_(~valid, -1.0e9)
 
-                selected_attention_key = compressed_key[batch_indices, safe_indices].float()
+                selected_attention_key = _gather_keys(
+                    compressed_key, batch_indices, safe_indices).float()
                 attention_scores = torch.einsum(
                     "bhcd,bckd->bhck",
                     attention_query[:, :, start:end].float(),

tests/ut/components/modules/test_compressed_key_gather.py

@@ -0,0 +1,135 @@
+# Copyright 2026 Huawei Technologies Co., Ltd
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Tests for gathering compressed keys in the CSA selection chains.
+
+The candidate chain and the Indexer KL loss gather compressed keys at selected
+positions. Advanced indexing with two index tensors has no AI Core
+implementation, so ``_gather_keys`` folds the batch offset into one index and
+uses a single ``index_select``. These tests pin that the substitution is the
+same gather, including for a batch larger than one, where the folded offset is
+what could go wrong.
+"""
+import torch
+
+import hyper_parallel.components.modules.shared_compressed_dsa_attention as csa
+from hyper_parallel.components.modules.shared_compressed_dsa_attention import (
+    _gather_keys,
+    compressed_candidate_topk,
+    compressed_causal_topk_and_candidates,
+)
+
+
+def _advanced_index(key: torch.Tensor, batch_indices: torch.Tensor,
+                    positions: torch.Tensor) -> torch.Tensor:
+    """The advanced indexing the helper replaces, kept as the reference."""
+    return key[batch_indices, positions]
+
+
+def _case(batch: int, compressed: int, rows: int, width: int, dim: int = 4):
+    """Key, batch ids and positions shaped the way the chains build them."""
+    torch.manual_seed(3)
+    key = torch.randn(batch, compressed, dim)
+    batch_indices = torch.arange(batch).view(-1, 1, 1)
+    positions = torch.randint(0, compressed, (batch, rows, width))
+    return key, batch_indices, positions
+
+
+def test_gather_matches_advanced_indexing_for_one_batch():
+    """The single-batch gather is unchanged."""
+    key, batch_indices, positions = _case(1, 16, 5, 7)
+    assert torch.equal(_gather_keys(key, batch_indices, positions),
+                       _advanced_index(key, batch_indices, positions))
+
+
+def test_gather_matches_advanced_indexing_across_batches():
+    """With batch > 1 the folded batch offset must pick the right rows."""
+    key, batch_indices, positions = _case(3, 16, 5, 7)
+    assert torch.equal(_gather_keys(key, batch_indices, positions),
+                       _advanced_index(key, batch_indices, positions))
+
+
+def test_gather_does_not_mix_batches():
+    """Each batch entry gathers only from its own key rows."""
+    key = torch.zeros(2, 4, 3)
+    key[1] = 1.0
+    batch_indices = torch.arange(2).view(-1, 1, 1)
+    positions = torch.zeros(2, 1, 2, dtype=torch.int64)
+    gathered = _gather_keys(key, batch_indices, positions)
+    assert torch.equal(gathered[0], torch.zeros(1, 2, 3))
+    assert torch.equal(gathered[1], torch.ones(1, 2, 3))
+
+
+def test_gather_shape_and_dtype():
+    """The result is positions.shape + (head_dim,) in the key's dtype."""
+    key, batch_indices, positions = _case(2, 8, 3, 5, dim=6)
+    gathered = _gather_keys(key, batch_indices, positions)
+    assert gathered.shape == (2, 3, 5, 6)
+    assert gathered.dtype == key.dtype
+
+
+def test_gather_handles_a_non_contiguous_key():
+    """A transposed key still gathers correctly."""
+    base = torch.randn(2, 4, 8)
+    key = base.transpose(1, 2)
+    batch_indices = torch.arange(2).view(-1, 1, 1)
+    positions = torch.randint(0, 8, (2, 3, 2))
+    assert torch.equal(_gather_keys(key, batch_indices, positions),
+                       _advanced_index(key, batch_indices, positions))
+
+
+def test_gather_handles_repeated_positions():
+    """Repeats are allowed; each occurrence gathers the same row."""
+    key, batch_indices, _ = _case(1, 8, 1, 1)
+    positions = torch.full((1, 2, 3), 5, dtype=torch.int64)
+    gathered = _gather_keys(key, batch_indices, positions)
+    assert torch.equal(gathered, key[0, 5].expand(1, 2, 3, key.shape[-1]))
+
+
+def _chain_inputs(seq: int = 16, comp_ratio: int = 2, heads: int = 2, dim: int = 4,
+                  batch: int = 1):
+    """Small CPU tensors shaped the way the chains expect."""
+    torch.manual_seed(5)
+    query = torch.randn(batch, seq, heads, dim)
+    key = torch.randn(batch, seq // comp_ratio, dim)
+    merge_weight = torch.randn(batch, seq, heads)
+    return query, key, merge_weight
+
+
+def test_candidate_topk_is_unchanged_by_the_gather(monkeypatch):
+    """The reindex chain selects the same positions either way."""
+    query, key, merge_weight = _chain_inputs()
+    candidates = compressed_causal_topk_and_candidates(
+        query, key, merge_weight, compress_ratio=2, sparse_count=4, topk_blocks=2,
+        block_size=2, query_chunk_size=8)[1]
+    args = {"compress_ratio": 2, "sparse_count": 4, "block_size": 2,
+            "query_chunk_size": 8}
+    fixed = compressed_candidate_topk(query, key, merge_weight, candidates, **args)
+    monkeypatch.setattr(csa, "_gather_keys", _advanced_index)
+    before = compressed_candidate_topk(query, key, merge_weight, candidates, **args)
+    assert torch.equal(fixed, before)
+
+
+def test_candidate_topk_is_unchanged_with_a_larger_batch(monkeypatch):
+    """Batch > 1 exercises the folded offset through the whole chain."""
+    query, key, merge_weight = _chain_inputs(batch=3)
+    candidates = compressed_causal_topk_and_candidates(
+        query, key, merge_weight, compress_ratio=2, sparse_count=4, topk_blocks=2,
+        block_size=2, query_chunk_size=8)[1]
+    args = {"compress_ratio": 2, "sparse_count": 4, "block_size": 2,
+            "query_chunk_size": 8}
+    fixed = compressed_candidate_topk(query, key, merge_weight, candidates, **args)
+    monkeypatch.setattr(csa, "_gather_keys", _advanced_index)
+    before = compressed_candidate_topk(query, key, merge_weight, candidates, **args)
+    assert torch.equal(fixed, before)