From ce675d3a74e7b8fef83fa686caadd1d5ea76f8c5 Mon Sep 17 00:00:00 2001 From: weich97 Date: Sat, 10 Oct 2026 23:53:07 -0400 Subject: [PATCH] perf(components): gather compressed keys with index_select, not two index tensors The candidate chain and the Indexer KL loss gather compressed keys at selected positions by indexing with two index tensors, a batch id and a position. That form of advanced indexing has no AI Core implementation, so it runs on AI CPU and then accounts for most of the time in both paths. A profiling breakdown of `compressed_candidate_topk` on the torch path (Ascend 910, 16 index heads, head_dim 128, bf16, sequence 4096, compressed length 2048) put 97.8% of device time on AI CPU, and named `Index` as 85.5% of it: 383 ms of a 449 ms step, with the position sort second at 50 ms. Folding the batch offset into the position index turns the gather into one `index_select` over the flattened key, which runs on device. Measured on the same card, with the gathered keys identical bit for bit: | gather shape | two tensors | index_select | ratio | |----------------------------------|-------------|--------------|--------| | 1024 rows x 1024 positions | 101.04 ms | 0.719 ms | 140.6x | | 1024 rows x 512 positions | 51.52 ms | 0.351 ms | 146.9x | | 1024 rows x 512, comp_len 16384 | 59.28 ms | 0.350 ms | 169.2x | End to end on the chain, with the selected positions identical bit for bit: `compressed_candidate_topk` at sequence 4096 and compressed length 2048 goes from 406.22 ms to 55.91 ms, a 350.30 ms saving per call. The first shape is what the candidate chain gathers per query chunk (`topk_blocks * block_size` positions); the second is what the KL loss gathers twice per chunk (`sparse_count` positions), once for the index key and once for the attention key. Both paths run per Indexer source layer per step, and the KL loss runs whenever the post-training Indexer stage is exercised. The gather is extracted into one helper so the three call sites stay consistent and the reason for the form is stated once. Unit tests cover equality with the advanced indexing for one batch and for several, that batches do not mix, the result shape and dtype, a non-contiguous key, repeated positions, and the unchanged output of the candidate chain at batch 1 and batch 3 (CPU only). --- .../shared_compressed_dsa_attention.py | 34 ++++- .../modules/test_compressed_key_gather.py | 135 ++++++++++++++++++ 2 files changed, 165 insertions(+), 4 deletions(-) create mode 100644 tests/ut/components/modules/test_compressed_key_gather.py diff --git a/hyper_parallel/components/modules/shared_compressed_dsa_attention.py b/hyper_parallel/components/modules/shared_compressed_dsa_attention.py index 5cb22cbdc..3cac45a44 100644 --- a/hyper_parallel/components/modules/shared_compressed_dsa_attention.py +++ b/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(), diff --git a/tests/ut/components/modules/test_compressed_key_gather.py b/tests/ut/components/modules/test_compressed_key_gather.py new file mode 100644 index 000000000..b9430586d --- /dev/null +++ b/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)